diff --git a/cohere/__init__.py b/cohere/__init__.py index 77e1f45fc..aa3548353 100644 --- a/cohere/__init__.py +++ b/cohere/__init__.py @@ -1,9 +1,5 @@ -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" @@ -11,134 +7,3 @@ 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 "" - return msg - - def __repr__(self) -> str: - return "%s(message=%r, http_status=%r, request_id=%r)" % ( - self.__class__.__name__, - self.message, - self.http_status, - ) - diff --git a/cohere/best_choices.py b/cohere/best_choices.py new file mode 100644 index 000000000..5e58c4b88 --- /dev/null +++ b/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) diff --git a/cohere/client.py b/cohere/client.py new file mode 100644 index 000000000..07a054a55 --- /dev/null +++ b/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 diff --git a/cohere/embeddings.py b/cohere/embeddings.py new file mode 100644 index 000000000..565189e1a --- /dev/null +++ b/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) diff --git a/cohere/error.py b/cohere/error.py new file mode 100644 index 000000000..3f56e5ea8 --- /dev/null +++ b/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 "" + return msg + + def __repr__(self) -> str: + return "%s(message=%r, http_status=%r, request_id=%r)" % ( + self.__class__.__name__, + self.message, + self.http_status, + ) diff --git a/cohere/generation.py b/cohere/generation.py new file mode 100644 index 000000000..e8f7d6a2e --- /dev/null +++ b/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 diff --git a/cohere/likelihoods.py b/cohere/likelihoods.py new file mode 100644 index 000000000..a9db3f632 --- /dev/null +++ b/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) diff --git a/cohere/similarities.py b/cohere/similarities.py new file mode 100644 index 000000000..6d2f25660 --- /dev/null +++ b/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) diff --git a/sanity-test.py b/sanity-test.py index 885dfffe3..bd61fcefc 100644 --- a/sanity-test.py +++ b/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: