Make co.generate call async (cohere-ai#74)

* Make co.generate async

* Add batched execution for generate, tokenize, and detokenize

* fix formatting

* formatting

* remove unused imports

* add tests

* fix tests

* fix bugs

* improve generation object

* convert simple objects into NamedTuples

* remove strange one-time iterator thing

* fix bugs

* remove unused import

* Apply suggestions from code review

Co-authored-by: Daulet Zhanguzin <[email protected]>

* Update cohere/classify.py

Co-authored-by: Michael <[email protected]>

* address comments

* fix flake8

* Update settings.json

* fix flake8

* rename key to index

Co-authored-by: Daulet Zhanguzin <[email protected]>
Co-authored-by: Michael <[email protected]>

Author: 3 people

cohere/classify.py

@@ -1,37 +1,39 @@
-from cohere.response import CohereObject
-from typing import List, Dict
+from typing import Dict, List, NamedTuple
 
+from cohere.response import CohereObject
 
-class LabelPrediction(CohereObject):
-    def __init__(self, confidence: float) -> None:
-        self.confidence = confidence
+LabelPrediction = NamedTuple("LabelPrediction", [("confidence", float)])
+Example = NamedTuple("Example", [("text", str), ("label", str)])
 
 
 class Classification(CohereObject):
-    def __init__(self, input: str,
-                 prediction: str, confidence: float, labels: Dict[str, LabelPrediction]) -> None:
+
+    def __init__(self, input: str, prediction: str, confidence: float, labels: Dict[str, LabelPrediction]) -> None:
         self.input = input
         self.prediction = prediction
         self.confidence = confidence
         self.labels = labels
 
+    def __repr__(self) -> str:
+        return f"Classification<prediction: \"{self.prediction}\", confidence: {self.confidence}>"
+
 
 class Classifications(CohereObject):
+
     def __init__(self, classifications: List[Classification]) -> None:
         self.classifications = classifications
-        self.iterator = iter(classifications)
 
-    def __iter__(self) -> iter:
-        return self.iterator
+    def __repr__(self) -> str:
+        return self.classifications.__repr__()
 
-    def __next__(self) -> next:
-        return next(self.iterator)
+    def __str__(self) -> str:
+        return self.classifications.__str__()
+
+    def __iter__(self) -> iter:
+        return iter(self.classifications)
 
     def __len__(self) -> int:
         return len(self.classifications)
 
-
-class Example(CohereObject):
-    def __init__(self, text: str, label: str) -> None:
-        self.text = text
-        self.label = label
+    def __getitem__(self, index) -> Classification:
+        return self.classifications[index]

cohere/client.py

@@ -2,23 +2,24 @@
 import math
 import sys
 from concurrent.futures import ThreadPoolExecutor
-from typing import Any, List, Dict
+from typing import Any, Dict, List
 from urllib.parse import urljoin
 
 import requests
 from requests import Response
 
 import cohere
-from cohere.classify import Classification, Classifications, LabelPrediction
+from cohere.classify import Classification, Classifications
 from cohere.classify import Example as ClassifyExample
+from cohere.classify import LabelPrediction
+from cohere.detokenize import Detokenization
 from cohere.embeddings import Embeddings
 from cohere.error import CohereError
 from cohere.extract import Entity
 from cohere.extract import Example as ExtractExample
 from cohere.extract import Extraction, Extractions
-from cohere.generation import Generation, Generations, TokenLikelihood
+from cohere.generation import Generations
 from cohere.tokenize import Tokens
-from cohere.detokenize import Detokenization
 
 use_xhr_client = False
 try:
@@ -33,12 +34,13 @@ class Client:
     def __init__(self,
                  api_key: str,
                  version: str = None,
-                 num_workers: int = 8,
+                 num_workers: int = 64,
                  request_dict: dict = {},
                  check_api_key: bool = True) -> None:
         self.api_key = api_key
         self.api_url = cohere.COHERE_API_URL
         self.batch_size = cohere.COHERE_EMBED_BATCH_SIZE
+        self._executor = ThreadPoolExecutor(num_workers)
         self.num_workers = num_workers
         self.request_dict = request_dict
         if version is None:
@@ -78,23 +80,28 @@ def check_api_key(self) -> Response:
             raise CohereError(message=res['message'], http_status=response.status_code, headers=response.headers)
         return res
 
-    def generate(
-        self,
-        prompt: str = None,
-        model: str = None,
-        preset: str = None,
-        num_generations: int = 1,
-        max_tokens: int = None,
-        temperature: float = 1.0,
-        k: int = 0,
-        p: float = 0.75,
-        frequency_penalty: float = 0.0,
-        presence_penalty: float = 0.0,
-        stop_sequences: List[str] = None,
-        return_likelihoods: str = 'NONE',
-        truncate: str = None,
-        logit_bias: Dict[int, float] = {}
-    ) -> Generations:
+    def batch_generate(self, prompts: List[str], **kwargs) -> List[Generations]:
+        generations: List[Generations] = []
+        for prompt in prompts:
+            kwargs["prompt"] = prompt
+            generations.append(self.generate(**kwargs))
+        return generations
+
+    def generate(self,
+                 prompt: str = None,
+                 model: str = None,
+                 preset: str = None,
+                 num_generations: int = 1,
+                 max_tokens: int = None,
+                 temperature: float = 1.0,
+                 k: int = 0,
+                 p: float = 0.75,
+                 frequency_penalty: float = 0.0,
+                 presence_penalty: float = 0.0,
+                 stop_sequences: List[str] = None,
+                 return_likelihoods: str = 'NONE',
+                 truncate: str = None,
+                 logit_bias: Dict[int, float] = {}) -> Generations:
         json_body = json.dumps({
             'model': model,
             'prompt': prompt,
@@ -111,21 +118,8 @@ def generate(
             'truncate': truncate,
             'logit_bias': logit_bias,
         })
-        response = self.__request(json_body, cohere.GENERATE_URL)
-
-        generations: List[Generation] = []
-        for gen in response['generations']:
-            likelihood = None
-            token_likelihoods = None
-            if return_likelihoods == 'GENERATION' or return_likelihoods == 'ALL':
-                likelihood = gen['likelihood']
-            if 'token_likelihoods' in gen.keys():
-                token_likelihoods = []
-                for likelihoods in gen['token_likelihoods']:
-                    token_likelihood = likelihoods['likelihood'] if 'likelihood' in likelihoods.keys() else None
-                    token_likelihoods.append(TokenLikelihood(likelihoods['token'], token_likelihood))
-            generations.append(Generation(gen['text'], likelihood, token_likelihoods))
-        return Generations(generations, return_likelihoods)
+        response = self._executor.submit(self.__request, json_body, cohere.GENERATE_URL)
+        return Generations(return_likelihoods=return_likelihoods, _future=response)
 
     def embed(self, texts: List[str], model: str = None, truncate: str = 'NONE') -> Embeddings:
         responses = []
@@ -146,22 +140,19 @@ def embed(self, texts: List[str], model: str = None, truncate: str = 'NONE') ->
                 response = self.__request(json_body, cohere.EMBED_URL)
                 responses.append(response['embeddings'])
         else:
-            with ThreadPoolExecutor(max_workers=self.num_workers) as executor:
-                for i in executor.map(self.__request, json_bodys, embed_url_stacked):
-                    request_futures.append(i)
+            for i in self._executor.map(self.__request, json_bodys, embed_url_stacked):
+                request_futures.append(i)
             for result in request_futures:
                 responses.extend(result['embeddings'])
 
         return Embeddings(responses)
 
-    def classify(
-        self,
-        inputs: List[str] = [],
-        model: str = None,
-        preset: str = None,
-        examples: List[ClassifyExample] = [],
-        truncate: str = None
-    ) -> Classifications:
+    def classify(self,
+                 inputs: List[str] = [],
+                 model: str = None,
+                 preset: str = None,
+                 examples: List[ClassifyExample] = [],
+                 truncate: str = None) -> Classifications:
         examples_dicts: list[dict[str, str]] = []
         for example in examples:
             example_dict = {'text': example.text, 'label': example.label}
@@ -209,19 +200,23 @@ def unstable_extract(self, examples: List[ExtractExample], texts: List[str]) ->
 
         return Extractions(extractions)
 
+    def batch_tokenize(self, texts: List[str]) -> List[Tokens]:
+        return [self.tokenize(t) for t in texts]
+
     def tokenize(self, text: str) -> Tokens:
         json_body = json.dumps({
             'text': text,
         })
-        response = self.__request(json_body, cohere.TOKENIZE_URL)
-        return Tokens(response['tokens'], response['token_strings'])
+        return Tokens(_future=self._executor.submit(self.__request, json_body, cohere.TOKENIZE_URL))
+
+    def batch_detokenize(self, list_of_tokens: List[List[int]]) -> List[Detokenization]:
+        return [self.detokenize(t) for t in list_of_tokens]
 
     def detokenize(self, tokens: List[int]) -> Detokenization:
         json_body = json.dumps({
             'tokens': tokens,
         })
-        response = self.__request(json_body, cohere.DETOKENIZE_URL)
-        return Detokenization(response['text'])
+        return Detokenization(_future=self._executor.submit(self.__request, json_body, cohere.DETOKENIZE_URL))
 
     def __print_warning_msg(self, response: Response):
         if 'X-API-Warning' in response.headers:

cohere/detokenize.py

@@ -1,6 +1,20 @@
-from cohere.response import CohereObject
+from concurrent.futures import Future
+from typing import Optional
+
+from cohere.response import AsyncAttribute, CohereObject
 
 
 class Detokenization(CohereObject):
-    def __init__(self, text: str) -> None:
-        self.text = text
+
+    def __init__(self, text: Optional[str] = None, *, _future: Optional[Future] = None) -> None:
+        if _future is not None:
+            self._init_from_future(_future)
+        else:
+            assert text is not None
+            self.text = text
+
+    def _init_from_future(self, future: Future):
+        self.text = AsyncAttribute(future, lambda x: x['text'])
+
+    def __str__(self) -> str:
+        return self.text

cohere/embeddings.py

@@ -1,15 +1,12 @@
+from typing import Iterator, List
+
 from cohere.response import CohereObject
-from typing import List
 
 
 class Embeddings(CohereObject):
 
     def __init__(self, embeddings: List[List[float]]) -> None:
         self.embeddings = embeddings
-        self.iterator = iter(embeddings)
-
-    def __iter__(self) -> iter:
-        return self.iterator
 
-    def __next__(self) -> next:
-        return next(self.iterator)
+    def __iter__(self) -> Iterator:
+        return iter(self.embeddings)

cohere/error.py

@@ -1,4 +1,5 @@
 class CohereError(Exception):
+
     def __init__(
         self,
         message=None,

cohere/extract.py

@@ -1,6 +1,7 @@
-from cohere.response import CohereObject
 from typing import List
 
+from cohere.response import CohereObject
+
 
 class Entity:
     '''
@@ -89,5 +90,5 @@ def __next__(self) -> next:
     def __len__(self) -> int:
         return len(self.extractions)
 
-    def __getitem__(self, index: int) -> Extraction:
+    def __getitem__(self, index) -> Extraction:
         return self.extractions[index]

cohere/generation.py

@@ -1,33 +1,70 @@
-from cohere.response import CohereObject
-from typing import List
+from concurrent.futures import Future
+from typing import Any, Dict, Iterator, List, NamedTuple, Optional, Union
 
+from cohere.response import AsyncAttribute, CohereObject
 
-class TokenLikelihood(CohereObject):
-    def __init__(self, token: str, likelihood: float) -> None:
-        self.token = token
-        self.likelihood = likelihood
+TokenLikelihood = NamedTuple("TokenLikelihood", [("token", str), ("likelihood", float)])
 
 
-class Generation(CohereObject):
-    def __init__(self,
-                 text: str,
-                 likelihood: float,
-                 token_likelihoods: List[TokenLikelihood]) -> None:
+class Generation(CohereObject, str):
+
+    def __new__(cls, text: str, *_, **__):
+        return str.__new__(cls, text)
+
+    def __init__(self, text: str, likelihood: float, token_likelihoods: List[TokenLikelihood]) -> None:
         self.text = text
         self.likelihood = likelihood
         self.token_likelihoods = token_likelihoods
 
+    def __str__(self) -> str:
+        return str(self.text)
+
+    def __len__(self) -> int:
+        return len(self.text)
+
+    def __getitem__(self, index) -> str:
+        return self.text[index]
+
 
 class Generations(CohereObject):
+
     def __init__(self,
-                 generations: List[Generation],
-                 return_likelihoods: str) -> None:
-        self.generations = generations
+                 return_likelihoods: str,
+                 response: Optional[Dict[str, Any]] = None,
+                 *,
+                 _future: Optional[Future] = None) -> None:
+        self.generations: Union[AsyncAttribute, List[Generation]] = None
         self.return_likelihoods = return_likelihoods
-        self.iterator = iter(generations)
+        if _future is not None:
+            self._init_from_future(_future)
+        else:
+            assert response is not None
+            self.generations = self._generations(response)
+
+    def _init_from_future(self, future: Future):
+        self.generations = AsyncAttribute(future, self._generations)
+
+    def _generations(self, response: Dict[str, Any]) -> List[Generation]:
+        generations: List[Generation] = []
+        for gen in response['generations']:
+            likelihood = None
+            token_likelihoods = None
+            if self.return_likelihoods in ['GENERATION', 'ALL']:
+                likelihood = gen['likelihood']
+            if 'token_likelihoods' in gen.keys():
+                token_likelihoods = []
+                for likelihoods in gen['token_likelihoods']:
+                    token_likelihood = likelihoods['likelihood'] if 'likelihood' in likelihoods.keys() else None
+                    token_likelihoods.append(TokenLikelihood(likelihoods['token'], token_likelihood))
+            generations.append(Generation(gen['text'], likelihood, token_likelihoods))
+
+        return generations
+
+    def __str__(self) -> str:
+        return str(self.generations)
 
-    def __iter__(self) -> iter:
-        return self.iterator
+    def __iter__(self) -> Iterator:
+        return iter(self.generations)
 
-    def __next__(self) -> next:
-        return next(self.iterator)
+    def __getitem__(self, index) -> Generation:
+        return self.generations[index]