Seperate files from __init__.py

cohere/__init__.py

@@ -1,144 +1,9 @@
-import json
-from typing import List
-from urllib.parse import urljoin
-
-import requests
-from requests import Response
+from .client import CohereClient
+from .error import CohereError
 
 COHERE_API_URL = "https://api.cohere.ai"
 GENERATE_URL = "generate"
 SIMILARITY_URL = "similarity"
 EMBED_URL = "embed"
 CHOOSE_BEST_URL = "choose-best"
 LIKELIHOOD_URL = "likelihood"
-
-class Generation:
-    def __init__(self, text) -> None:
-        self.text = text
-
-    def __str__(self) -> str:
-        return self.text
-
-class Similarities:
-    def __init__(self, similarities) -> None:
-        self.similarities = similarities
-
-    def __str__(self) -> str:
-        return str(self.similarities)
-
-class Embeddings:
-    def __init__(self, embeddings) -> None:
-        self.embeddings = embeddings
-
-    def __str__(self) -> str:
-        return str(self.embeddings)
-
-class BestChoices:
-    def __init__(self, likelihoods, mode) -> None:
-        self.likelihoods = likelihoods
-        self.mode = mode
-
-    def __str__(self) -> str:
-        return str(self.likelihoods)
-
-class Likelihoods:
-    def __init__(self, likelihood, token_likelihoods) -> None:
-        self.likelihood = likelihood
-        self.token_likelihoods = token_likelihoods
-
-    def __str__(self) -> str:
-        return str(self.likelihood) + "\n" + str(self.token_likelihoods)
-
-class CohereClient:
-    def __init__(self, api_key: str) -> None:
-        self.api_key = api_key
-        self.api_url = COHERE_API_URL
-
-    def generate(
-        self, 
-        model: str, 
-        prompt: str, 
-        max_tokens: int = 20, 
-        temperature: float = 1.0, 
-        k: int = 0, p: float = 0.75
-    ) -> Generation:
-        json_body = json.dumps({
-            "prompt": prompt,
-            "max_tokens": max_tokens,
-            "temperature": temperature,
-            "k": k,
-            "p": p,
-        })
-        response = self.__request(json_body, GENERATE_URL, model)
-        return Generation(response["text"])
-
-    def similarity(self, model: str, anchor: str, targets: List[str]) -> Similarities:
-        json_body = json.dumps({
-            "anchor": anchor,
-            "targets": targets,
-        })
-        response = self.__request(json_body, SIMILARITY_URL, model)
-        return Similarities(response["similarities"])
-
-    def embed(self, model: str, texts: List[str]) -> Embeddings:
-        json_body = json.dumps({
-            "texts": texts,
-        })
-        response = self.__request(json_body, EMBED_URL, model)
-        return Embeddings(response["embeddings"])
-
-    def choose_best(self, model: str, query: str, options: List[str], mode:  str = "") -> BestChoices:
-        json_body = json.dumps({
-            "query": query,
-            "options": options,
-            "mode": mode,
-        })
-        response = self.__request(json_body, CHOOSE_BEST_URL, model)
-        return BestChoices(response['likelihoods'], mode)
-
-    def likelihood(self, model: str, text: List[str]) -> Likelihoods:
-        json_body = json.dumps({
-            "text": text,
-        })
-        response = self.__request(json_body, LIKELIHOOD_URL, model)
-        return Likelihoods(response['likelihood'], response['token_likelihoods'])
-
-    def __request(self, json_body, endpoint, model) -> Response:
-        headers = {
-            'Authorization': 'BEARER {}'.format(self.api_key),
-            'Content-Type': 'application/json'
-        }
-        url = urljoin(self.api_url, model + "/" + endpoint)
-        response = requests.request("POST", url, headers=headers, data=json_body)
-        res = json.loads(response.text)
-        if "message" in res.keys(): # has errors
-            raise CohereError(
-                message=res["message"],
-                http_status=response.status_code,
-                headers=response.headers)
-        return res
-
-class CohereError(Exception):
-    def __init__(
-        self,
-        message=None,
-        http_status=None,
-        headers=None,
-    ) -> None:
-        super(CohereError, self).__init__(message)
-
-        self.message = message
-        self.http_status = http_status
-        self.headers = headers or {}
-
-    def __str__(self) -> str:
-        msg = self.message or "<empty message>"
-        return msg
-
-    def __repr__(self) -> str:
-        return "%s(message=%r, http_status=%r, request_id=%r)" % (
-            self.__class__.__name__,
-            self.message,
-            self.http_status,
-        )
-

cohere/best_choices.py

@@ -0,0 +1,9 @@
+from typing import List
+
+class BestChoices:
+    def __init__(self, likelihoods: List[float], mode: str) -> None:
+        self.likelihoods = likelihoods
+        self.mode = mode
+
+    def __str__(self) -> str:
+        return str(self.likelihoods)

cohere/client.py

@@ -0,0 +1,85 @@
+import json
+from urllib.parse import urljoin
+
+from typing import List
+import requests
+from requests import Response
+
+import cohere
+from cohere.error import CohereError
+
+from cohere.generation import Generation
+from cohere.similarities import Similarities
+from cohere.embeddings import Embeddings
+from cohere.best_choices import BestChoices
+from cohere.likelihoods import Likelihoods
+
+
+class CohereClient:
+    def __init__(self, api_key: str) -> None:
+        self.api_key = api_key
+        self.api_url = cohere.COHERE_API_URL
+
+    def generate(
+        self, 
+        model: str, 
+        prompt: str, 
+        max_tokens: int = 20, 
+        temperature: float = 1.0, 
+        k: int = 0, p: float = 0.75
+    ) -> Generation:
+        json_body = json.dumps({
+            "prompt": prompt,
+            "max_tokens": max_tokens,
+            "temperature": temperature,
+            "k": k,
+            "p": p,
+        })
+        response = self.__request(json_body, cohere.GENERATE_URL, model)
+        return Generation(response["text"])
+
+    def similarity(self, model: str, anchor: str, targets: List[str]) -> Similarities:
+        json_body = json.dumps({
+            "anchor": anchor,
+            "targets": targets,
+        })
+        response = self.__request(json_body, cohere.SIMILARITY_URL, model)
+        return Similarities(response["similarities"])
+
+    def embed(self, model: str, texts: List[str]) -> Embeddings:
+        json_body = json.dumps({
+            "texts": texts,
+        })
+        response = self.__request(json_body, cohere.EMBED_URL, model)
+        return Embeddings(response["embeddings"])
+
+    def choose_best(self, model: str, query: str, options: List[str], mode:  str = "") -> BestChoices:
+        json_body = json.dumps({
+            "query": query,
+            "options": options,
+            "mode": mode,
+        })
+        response = self.__request(json_body, cohere.CHOOSE_BEST_URL, model)
+        return BestChoices(response['likelihoods'], mode)
+
+    def likelihood(self, model: str, text: List[str]) -> Likelihoods:
+        json_body = json.dumps({
+            "text": text,
+        })
+        response = self.__request(json_body, cohere.LIKELIHOOD_URL, model)
+        return Likelihoods(response['likelihood'], response['token_likelihoods'])
+
+    def __request(self, json_body, endpoint, model) -> Response:
+        headers = {
+            'Authorization': 'BEARER {}'.format(self.api_key),
+            'Content-Type': 'application/json'
+        }
+        url = urljoin(self.api_url, model + "/" + endpoint)
+        response = requests.request("POST", url, headers=headers, data=json_body)
+        res = json.loads(response.text)
+        if "message" in res.keys(): # has errors
+            raise CohereError(
+                message=res["message"],
+                http_status=response.status_code,
+                headers=response.headers)
+        return res

cohere/embeddings.py

@@ -0,0 +1,8 @@
+from typing import List
+
+class Embeddings:
+    def __init__(self, embeddings: List[List[float]]) -> None:
+        self.embeddings = embeddings
+
+    def __str__(self) -> str:
+        return str(self.embeddings)

cohere/error.py

@@ -0,0 +1,23 @@
+class CohereError(Exception):
+    def __init__(
+        self,
+        message=None,
+        http_status=None,
+        headers=None,
+    ) -> None:
+        super(CohereError, self).__init__(message)
+
+        self.message = message
+        self.http_status = http_status
+        self.headers = headers or {}
+
+    def __str__(self) -> str:
+        msg = self.message or "<empty message>"
+        return msg
+
+    def __repr__(self) -> str:
+        return "%s(message=%r, http_status=%r, request_id=%r)" % (
+            self.__class__.__name__,
+            self.message,
+            self.http_status,
+        )

cohere/generation.py

@@ -0,0 +1,6 @@
+class Generation:
+    def __init__(self, text: str) -> None:
+        self.text = text
+
+    def __str__(self) -> str:
+        return self.text

cohere/likelihoods.py

@@ -0,0 +1,9 @@
+from typing import List, Dict
+
+class Likelihoods:
+    def __init__(self, likelihood: float, token_likelihoods: List[Dict]) -> None:
+        self.likelihood = likelihood
+        self.token_likelihoods = token_likelihoods
+
+    def __str__(self) -> str:
+        return str(self.likelihood) + "\n" + str(self.token_likelihoods)

cohere/similarities.py

@@ -0,0 +1,8 @@
+from typing import List
+
+class Similarities:
+    def __init__(self, similarities: List[float]) -> None:
+        self.similarities = similarities
+
+    def __str__(self) -> str:
+        return str(self.similarities)

sanity-test.py

@@ -7,30 +7,31 @@
 co = cohere.CohereClient(API_KEY)
 
 prediction = co.generate(
-            model="baseline-355m",
+            model="baseline-shrimp",
             prompt="co:here",
             max_tokens=10)
 print('prediction: {}'.format(prediction.text))
 
 embeddings = co.embed(
-            model="baseline-124m",
+            model="baseline-shrimp",
             texts=["co:here", "cohere"])
 
 similarities = co.similarity(
-            model="baseline-124m",
+            model="baseline-shrimp",
            	anchor="cohere ai",
             targets=["co:here", "cohere"])
 print('Similarity value of `co:here`: {}'.format(similarities.similarities[0]))
 
 options = co.choose_best(
-            model="baseline-355m",
+            model="baseline-shrimp",
             query="hello {}",
-            options=["world", "cohere"])
+            options=["world", "cohere"],
+            mode="APPEND_OPTION")
 print('first option is `world`, with likelihood value of {}'.format(options.likelihoods[0]))
-print('Selected mode was {}'.format(options['mode']))
+print('Selected mode was {}'.format(options.mode))
 
 likelihood = co.likelihood(
-            model="baseline-355m",
+            model="baseline-shrimp",
             text="hello, my name is johnny SURPRISE")
 print('likelihood of text is {}'.format(likelihood.likelihood))
 print('token likelihoods are: (first token has no likelihood)')
@@ -46,7 +47,7 @@
 
 try:
 	predictions = co.generate(
-            model="baseline-124m",
+            model="baseline-shrimp",
             prompt="",
             max_tokens=10)
 except cohere.CohereError as e: