Updated documentation

README.md

@@ -7,6 +7,10 @@ Transformer-based zero/few shot learning components for scikit-learn pipelines.
 
 [Documentation](https://centre-for-humanities-computing.github.io/stormtrooper/)
 
+## New in version 0.4.0 :fire:
+
+- You can now use OpenAI's chat models with blazing fast :zap: async inference.
+
 ## New in version 0.3.0 🌟 
 
 - SetFit is now part of the library and can be used in scikit-learn workflows.
@@ -71,9 +75,24 @@ predictions = classifier.predict(example_texts)
 assert list(predictions) == ["atheism/christianity", "astronomy/space"]
 ```
 
+OpenAI models:
+You can now use OpenAI's chat LLMs in stormtrooper workflows.
+
+```python
+from stormtrooper import OpenAIZeroShotClassifier
+
+classifier = OpenAIZeroShotClassifier("gpt-4").fit(None, class_labels)
+```
+
+```python
+predictions = classifier.predict(example_texts)
+
+assert list(predictions) == ["atheism/christianity", "astronomy/space"]
+```
+
 ### Few-Shot Learning
 
-For few-shot tasks you can only use Generative, Text2Text (aka. promptable) or SetFit models.
+For few-shot tasks you can only use Generative, Text2Text, OpenAI (aka. promptable) or SetFit models.
 
 ```python
 from stormtrooper import GenerativeFewShotClassifier, Text2TextFewShotClassifier, SetFitFewShotClassifier

docs/_build/html/_sources/index.rst.txt

@@ -17,6 +17,14 @@ If you intend to use SetFit models as well, install stormtrooper with optional d
 
    pip install stormtrooper[setfit]
 
+From version 0.4.0 you can also use OpenAI models in stormtrooper.
+
+.. code-block::
+
+   pip install stormtrooper[openai]
+   export OPENAI_API_KEY="sk-..."
+
+
 Usage
 ^^^^^^^^^
 
@@ -45,6 +53,7 @@ In this example I am going to use Google's FLAN-T5.
    text2text
    generative
    setfit
+   openai
    prompting
    inference_on_gpu
 

docs/_build/html/_sources/openai.rst.txt

@@ -0,0 +1,85 @@
+OpenAI models
+=================
+
+Stormtrooper gives you access to OpenAI's chat models for zero and few-shot classification.
+You get full control over temperature settings, system and user prompts.
+In contrast to other packages, like scikit-llm, stormtrooper also uses Python's asyncio to concurrently
+interact with OpenAI's API. This can give multiple times speedup on several tasks. 
+
+You can also set upper limits for number of requests and tokens per minute, so you don't exceed your quota.
+This is by default set to the limit of the payed tier on OpenAI's API.
+
+You need to install stormtrooper with optional dependencies.
+
+.. code-block:: bash
+
+   pip install stormtrooper[openai]
+
+You additionally need to set the OpenAI API key as an environment variable.
+
+.. code-block:: bash
+
+   export OPENAI_API_KEY="sk-..."
+   # Setting organization is optional
+   export OPENAI_ORG="org-..."
+
+.. code-block:: python
+
+   from stormtrooper import OpenAIZeroShotClassifier, OpenAIFewShotClassifier
+
+   sample_text = "It is the Electoral College's responsibility to elect the president."
+
+   labels = ["politics", "science", "other"]
+
+Here's a zero shot example with ChatGPT 3.5:
+
+.. code-block:: python
+
+   model = OpenAIZeroShotClassifier("gpt-3.5-turbo").fit(None, labels)
+   predictions = model.predict([sample_text])
+   assert list(predictions) == ["politics"]
+   
+And a few shot example with ChatGPT 4:
+
+.. code-block:: python
+
+   few_shot_examples = [
+     "Joe Biden is the president.",
+     "Liquid water was found on the moon.",
+     "Jerry likes football."
+   ]
+
+   model = OpenAIFewShotClassifier("gpt-4", temperature=0.2).fit(few_shot_examples, labels)
+   predictions = model.predict([sample_text])
+   assert list(predictions) == ["politics"]
+
+
+The format of the prompts is the same as with StableBeluga instruct models, and an error is raised if your prompt does not follow
+this format.
+
+.. code-block:: python
+  
+   prompt = """
+   ### System:
+   You are a helpful assistant
+   ### User:
+   Your task will be to classify a text document into one
+   of the following classes: {classes}.
+   Please respond with a single label that you think fits
+   the document best.
+   Classify the following piece of text:
+   '{X}'
+   ### Assistant:
+   """
+
+   model = OpenAIZeroShotClassifier("gpt-4", prompt=prompt)
+
+
+API reference
+^^^^^^^^^^^^^
+
+.. autoclass:: stormtrooper.OpenAIZeroShotClassifier
+   :members:
+
+.. autoclass:: stormtrooper.OpenAIFewShotClassifier
+   :members:

docs/_build/html/_sources/prompting.rst.txt

@@ -1,7 +1,7 @@
 Prompting
 =========
 
-Text2Text and Generative models use a prompting approach for classification.
+Text2Text, Generative, and OpenAI models use a prompting approach for classification.
 stormtrooper comes with default prompts, but these might not suit the model you want to use,
 or your use case might require a different prompting strategy from the default.
 stormtrooper allows you to specify custom prompts in these cases.
@@ -12,7 +12,7 @@ Templates
 Prompting in stormtrooper uses a templating approach, where the .format() method is called on prompts to
 insert labels and data.
 
-A zero-shot prompt for an instruct Llama model like Stable Beluga would look something like this (this is the default):
+A zero-shot prompt for an instruct Llama model like Stable Beluga or for ChatGPT would look something like this (this is the default):
 
 .. code-block:: python
   

docs/_build/html/generative.html

@@ -5,7 +5,7 @@
     <meta name="color-scheme" content="light dark"><meta name="viewport" content="width=device-width, initial-scale=1" />
 <link rel="index" title="Index" href="genindex.html" /><link rel="search" title="Search" href="search.html" /><link rel="next" title="SetFit models" href="setfit.html" /><link rel="prev" title="Text2Text models" href="text2text.html" />
 
-    <link rel="shortcut icon" href="_static/logo.svg"/><!-- Generated with Sphinx 7.2.5 and Furo 2023.09.10 -->
+    <link rel="shortcut icon" href="_static/logo.svg"/><!-- Generated with Sphinx 7.2.6 and Furo 2023.09.10 -->
         <title>Generative models - stormtrooper</title>
       <link rel="stylesheet" type="text/css" href="_static/pygments.css?v=362ab14a" />
     <link rel="stylesheet" type="text/css" href="_static/styles/furo.css?v=135e06be" />
@@ -170,6 +170,7 @@
 <li class="toctree-l1"><a class="reference internal" href="text2text.html">Text2Text models</a></li>
 <li class="toctree-l1 current current-page"><a class="current reference internal" href="#">Generative models</a></li>
 <li class="toctree-l1"><a class="reference internal" href="setfit.html">SetFit models</a></li>
+<li class="toctree-l1"><a class="reference internal" href="openai.html">OpenAI models</a></li>
 <li class="toctree-l1"><a class="reference internal" href="prompting.html">Prompting</a></li>
 <li class="toctree-l1"><a class="reference internal" href="inference_on_gpu.html">Inference on GPU</a></li>
 </ul>

docs/_build/html/genindex.html

@@ -4,7 +4,7 @@
     <meta name="viewport" content="width=device-width,initial-scale=1"/>
     <meta name="color-scheme" content="light dark"><link rel="index" title="Index" href="#" /><link rel="search" title="Search" href="search.html" />
 
-    <link rel="shortcut icon" href="_static/logo.svg"/><!-- Generated with Sphinx 7.2.5 and Furo 2023.09.10 --><title>Index - stormtrooper</title>
+    <link rel="shortcut icon" href="_static/logo.svg"/><!-- Generated with Sphinx 7.2.6 and Furo 2023.09.10 --><title>Index - stormtrooper</title>
 <link rel="stylesheet" type="text/css" href="_static/pygments.css?v=362ab14a" />
     <link rel="stylesheet" type="text/css" href="_static/styles/furo.css?v=135e06be" />
     <link rel="stylesheet" type="text/css" href="_static/styles/furo-extensions.css?v=36a5483c" />
@@ -168,6 +168,7 @@
 <li class="toctree-l1"><a class="reference internal" href="text2text.html">Text2Text models</a></li>
 <li class="toctree-l1"><a class="reference internal" href="generative.html">Generative models</a></li>
 <li class="toctree-l1"><a class="reference internal" href="setfit.html">SetFit models</a></li>
+<li class="toctree-l1"><a class="reference internal" href="openai.html">OpenAI models</a></li>
 <li class="toctree-l1"><a class="reference internal" href="prompting.html">Prompting</a></li>
 <li class="toctree-l1"><a class="reference internal" href="inference_on_gpu.html">Inference on GPU</a></li>
 </ul>
@@ -209,7 +210,7 @@
 
 <section class="genindex-section">
   <h1 id="index">Index</h1>
-  <div class="genindex-jumpbox"><a href="#C"><strong>C</strong></a> | <a href="#E"><strong>E</strong></a> | <a href="#F"><strong>F</strong></a> | <a href="#G"><strong>G</strong></a> | <a href="#P"><strong>P</strong></a> | <a href="#R"><strong>R</strong></a> | <a href="#S"><strong>S</strong></a> | <a href="#T"><strong>T</strong></a> | <a href="#Z"><strong>Z</strong></a></div>
+  <div class="genindex-jumpbox"><a href="#C"><strong>C</strong></a> | <a href="#E"><strong>E</strong></a> | <a href="#F"><strong>F</strong></a> | <a href="#G"><strong>G</strong></a> | <a href="#O"><strong>O</strong></a> | <a href="#P"><strong>P</strong></a> | <a href="#R"><strong>R</strong></a> | <a href="#S"><strong>S</strong></a> | <a href="#T"><strong>T</strong></a> | <a href="#Z"><strong>Z</strong></a></div>
 </section>
 <section id="C" class="genindex-section">
   <h2>C</h2>
@@ -219,6 +220,10 @@ <h2>C</h2>
 
         <ul>
           <li><a href="generative.html#stormtrooper.GenerativeZeroShotClassifier.classes_">(stormtrooper.GenerativeZeroShotClassifier attribute)</a>
+</li>
+          <li><a href="openai.html#stormtrooper.OpenAIFewShotClassifier.classes_">(stormtrooper.OpenAIFewShotClassifier attribute)</a>
+</li>
+          <li><a href="openai.html#stormtrooper.OpenAIZeroShotClassifier.classes_">(stormtrooper.OpenAIZeroShotClassifier attribute)</a>
 </li>
           <li><a href="setfit.html#stormtrooper.SetFitFewShotClassifier.classes_">(stormtrooper.SetFitFewShotClassifier attribute)</a>
 </li>
@@ -257,6 +262,10 @@ <h2>F</h2>
 
         <ul>
           <li><a href="generative.html#stormtrooper.GenerativeZeroShotClassifier.fit">(stormtrooper.GenerativeZeroShotClassifier method)</a>
+</li>
+          <li><a href="openai.html#stormtrooper.OpenAIFewShotClassifier.fit">(stormtrooper.OpenAIFewShotClassifier method)</a>
+</li>
+          <li><a href="openai.html#stormtrooper.OpenAIZeroShotClassifier.fit">(stormtrooper.OpenAIZeroShotClassifier method)</a>
 </li>
           <li><a href="setfit.html#stormtrooper.SetFitFewShotClassifier.fit">(stormtrooper.SetFitFewShotClassifier method)</a>
 </li>
@@ -293,6 +302,20 @@ <h2>G</h2>
   </tr></table>
 </section>
 
+<section id="O" class="genindex-section">
+  <h2>O</h2>
+  <table style="width: 100%" class="indextable genindextable"><tr>
+    <td style="width: 33%; vertical-align: top;"><ul>
+        <li><a href="openai.html#stormtrooper.OpenAIFewShotClassifier">OpenAIFewShotClassifier (class in stormtrooper)</a>
+</li>
+    </ul></td>
+    <td style="width: 33%; vertical-align: top;"><ul>
+        <li><a href="openai.html#stormtrooper.OpenAIZeroShotClassifier">OpenAIZeroShotClassifier (class in stormtrooper)</a>
+</li>
+    </ul></td>
+  </tr></table>
+</section>
+
 <section id="P" class="genindex-section">
   <h2>P</h2>
   <table style="width: 100%" class="indextable genindextable"><tr>
@@ -301,6 +324,10 @@ <h2>P</h2>
 
         <ul>
           <li><a href="generative.html#stormtrooper.GenerativeZeroShotClassifier.partial_fit">(stormtrooper.GenerativeZeroShotClassifier method)</a>
+</li>
+          <li><a href="openai.html#stormtrooper.OpenAIFewShotClassifier.partial_fit">(stormtrooper.OpenAIFewShotClassifier method)</a>
+</li>
+          <li><a href="openai.html#stormtrooper.OpenAIZeroShotClassifier.partial_fit">(stormtrooper.OpenAIZeroShotClassifier method)</a>
 </li>
           <li><a href="setfit.html#stormtrooper.SetFitZeroShotClassifier.partial_fit">(stormtrooper.SetFitZeroShotClassifier method)</a>
 </li>
@@ -315,6 +342,10 @@ <h2>P</h2>
 
         <ul>
           <li><a href="generative.html#stormtrooper.GenerativeZeroShotClassifier.predict">(stormtrooper.GenerativeZeroShotClassifier method)</a>
+</li>
+          <li><a href="openai.html#stormtrooper.OpenAIFewShotClassifier.predict">(stormtrooper.OpenAIFewShotClassifier method)</a>
+</li>
+          <li><a href="openai.html#stormtrooper.OpenAIZeroShotClassifier.predict">(stormtrooper.OpenAIZeroShotClassifier method)</a>
 </li>
           <li><a href="setfit.html#stormtrooper.SetFitFewShotClassifier.predict">(stormtrooper.SetFitFewShotClassifier method)</a>
 </li>
@@ -355,6 +386,10 @@ <h2>S</h2>
 
         <ul>
           <li><a href="generative.html#stormtrooper.GenerativeZeroShotClassifier.set_score_request">(stormtrooper.GenerativeZeroShotClassifier method)</a>
+</li>
+          <li><a href="openai.html#stormtrooper.OpenAIFewShotClassifier.set_score_request">(stormtrooper.OpenAIFewShotClassifier method)</a>
+</li>
+          <li><a href="openai.html#stormtrooper.OpenAIZeroShotClassifier.set_score_request">(stormtrooper.OpenAIZeroShotClassifier method)</a>
 </li>
           <li><a href="setfit.html#stormtrooper.SetFitFewShotClassifier.set_score_request">(stormtrooper.SetFitFewShotClassifier method)</a>
 </li>

docs/_build/html/index.html

@@ -5,7 +5,7 @@
     <meta name="color-scheme" content="light dark"><meta name="viewport" content="width=device-width, initial-scale=1" />
 <link rel="index" title="Index" href="genindex.html" /><link rel="search" title="Search" href="search.html" /><link rel="next" title="Zero-shot models" href="zeroshot.html" />
 
-    <link rel="shortcut icon" href="_static/logo.svg"/><!-- Generated with Sphinx 7.2.5 and Furo 2023.09.10 -->
+    <link rel="shortcut icon" href="_static/logo.svg"/><!-- Generated with Sphinx 7.2.6 and Furo 2023.09.10 -->
         <title>stormtrooper</title>
       <link rel="stylesheet" type="text/css" href="_static/pygments.css?v=362ab14a" />
     <link rel="stylesheet" type="text/css" href="_static/styles/furo.css?v=135e06be" />
@@ -170,6 +170,7 @@
 <li class="toctree-l1"><a class="reference internal" href="text2text.html">Text2Text models</a></li>
 <li class="toctree-l1"><a class="reference internal" href="generative.html">Generative models</a></li>
 <li class="toctree-l1"><a class="reference internal" href="setfit.html">SetFit models</a></li>
+<li class="toctree-l1"><a class="reference internal" href="openai.html">OpenAI models</a></li>
 <li class="toctree-l1"><a class="reference internal" href="prompting.html">Prompting</a></li>
 <li class="toctree-l1"><a class="reference internal" href="inference_on_gpu.html">Inference on GPU</a></li>
 </ul>
@@ -223,6 +224,11 @@ <h2>Installation<a class="headerlink" href="#installation" title="Link to this h
 <div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="n">pip</span> <span class="n">install</span> <span class="n">stormtrooper</span><span class="p">[</span><span class="n">setfit</span><span class="p">]</span>
 </pre></div>
 </div>
+<p>From version 0.4.0 you can also use OpenAI models in stormtrooper.</p>
+<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="n">pip</span> <span class="n">install</span> <span class="n">stormtrooper</span><span class="p">[</span><span class="n">openai</span><span class="p">]</span>
+<span class="n">export</span> <span class="n">OPENAI_API_KEY</span><span class="o">=</span><span class="s2">&quot;sk-...&quot;</span>
+</pre></div>
+</div>
 </section>
 <section id="usage">
 <h2>Usage<a class="headerlink" href="#usage" title="Link to this heading">#</a></h2>
@@ -247,6 +253,7 @@ <h2>Usage<a class="headerlink" href="#usage" title="Link to this heading">#</a><
 <li class="toctree-l1"><a class="reference internal" href="text2text.html">Text2Text models</a></li>
 <li class="toctree-l1"><a class="reference internal" href="generative.html">Generative models</a></li>
 <li class="toctree-l1"><a class="reference internal" href="setfit.html">SetFit models</a></li>
+<li class="toctree-l1"><a class="reference internal" href="openai.html">OpenAI models</a></li>
 <li class="toctree-l1"><a class="reference internal" href="prompting.html">Prompting</a></li>
 <li class="toctree-l1"><a class="reference internal" href="inference_on_gpu.html">Inference on GPU</a></li>
 </ul>

docs/_build/html/inference_on_gpu.html

@@ -5,7 +5,7 @@
     <meta name="color-scheme" content="light dark"><meta name="viewport" content="width=device-width, initial-scale=1" />
 <link rel="index" title="Index" href="genindex.html" /><link rel="search" title="Search" href="search.html" /><link rel="prev" title="Prompting" href="prompting.html" />
 
-    <link rel="shortcut icon" href="_static/logo.svg"/><!-- Generated with Sphinx 7.2.5 and Furo 2023.09.10 -->
+    <link rel="shortcut icon" href="_static/logo.svg"/><!-- Generated with Sphinx 7.2.6 and Furo 2023.09.10 -->
         <title>Inference on GPU - stormtrooper</title>
       <link rel="stylesheet" type="text/css" href="_static/pygments.css?v=362ab14a" />
     <link rel="stylesheet" type="text/css" href="_static/styles/furo.css?v=135e06be" />
@@ -170,6 +170,7 @@
 <li class="toctree-l1"><a class="reference internal" href="text2text.html">Text2Text models</a></li>
 <li class="toctree-l1"><a class="reference internal" href="generative.html">Generative models</a></li>
 <li class="toctree-l1"><a class="reference internal" href="setfit.html">SetFit models</a></li>
+<li class="toctree-l1"><a class="reference internal" href="openai.html">OpenAI models</a></li>
 <li class="toctree-l1"><a class="reference internal" href="prompting.html">Prompting</a></li>
 <li class="toctree-l1 current current-page"><a class="current reference internal" href="#">Inference on GPU</a></li>
 </ul>