From 5791bc0c972671d6d963f486f20b7c3979a6f7d8 Mon Sep 17 00:00:00 2001 From: Alekhya Date: Sun, 22 Oct 2023 01:50:30 -0700 Subject: [PATCH 1/6] add compression parameter to embed --- cohere/client.py | 9 +++------ cohere/client_async.py | 9 +++------ 2 files changed, 6 insertions(+), 12 deletions(-) diff --git a/cohere/client.py b/cohere/client.py index 266b250ed..5d671381b 100644 --- a/cohere/client.py +++ b/cohere/client.py @@ -393,8 +393,7 @@ def embed( texts: List[str], model: Optional[str] = None, truncate: Optional[str] = None, - compress: Optional[bool] = False, - compression_codebook: Optional[str] = "default", + compression: Optional[str] = None, input_type: Optional[str] = None, ) -> Embeddings: """Returns an Embeddings object for the provided texts. Visit https://cohere.ai/embed to learn about embeddings. @@ -403,8 +402,7 @@ def embed( text (List[str]): A list of strings to embed. model (str): (Optional) The model ID to use for embedding the text. truncate (str): (Optional) One of NONE|START|END, defaults to END. How the API handles text longer than the maximum token length. - compress (bool): (Optional) Whether to compress the embeddings. When True, the compressed_embeddings will be returned as integers in the range [0, 255]. - compression_codebook (str): (Optional) The compression codebook to use for compressed embeddings. Defaults to "default". + compression (str): (Optional) One of "int8" or "binary". The type of compression to use for the embeddings. input_type (str): (Optional) One of "classification", "clustering", "search_document", "search_query". The type of input text provided to embed. """ responses = { @@ -420,8 +418,7 @@ def embed( "model": model, "texts": texts_batch, "truncate": truncate, - "compress": compress, - "compression_codebook": compression_codebook, + "compression": compression, "input_type": input_type, } ) diff --git a/cohere/client_async.py b/cohere/client_async.py index bbc541c71..f67708caa 100644 --- a/cohere/client_async.py +++ b/cohere/client_async.py @@ -271,8 +271,7 @@ async def embed( texts: List[str], model: Optional[str] = None, truncate: Optional[str] = None, - compress: Optional[bool] = False, - compression_codebook: Optional[str] = "default", + compression: Optional[str] = None, input_type: Optional[str] = None, ) -> Embeddings: """Returns an Embeddings object for the provided texts. Visit https://cohere.ai/embed to learn about embeddings. @@ -281,8 +280,7 @@ async def embed( text (List[str]): A list of strings to embed. model (str): (Optional) The model ID to use for embedding the text. truncate (str): (Optional) One of NONE|START|END, defaults to END. How the API handles text longer than the maximum token length. - compress (bool): (Optional) Whether to compress the embeddings. When True, the compressed_embeddings will be returned as integers in the range [0, 255]. - compression_codebook (str): (Optional) The compression codebook to use for compressed embeddings. Defaults to "default". + compression (str): (Optional) One of "int8" or "binary". The type of compression to use for the embeddings. input_type (str): (Optional) One of "classification", "clustering", "search_document", "search_query". The type of input text provided to embed. """ json_bodys = [ @@ -290,8 +288,7 @@ async def embed( texts=texts[i : i + cohere.COHERE_EMBED_BATCH_SIZE], model=model, truncate=truncate, - compress=compress, - compression_codebook=compression_codebook, + compression=compression, input_type=input_type, ) for i in range(0, len(texts), cohere.COHERE_EMBED_BATCH_SIZE) From 7298db5184805ea2cc3fe1fc3d9445395c83dfe3 Mon Sep 17 00:00:00 2001 From: Alekhya Date: Sun, 22 Oct 2023 01:54:18 -0700 Subject: [PATCH 2/6] update changelog and toml --- CHANGELOG.md | 3 +++ pyproject.toml | 2 +- 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index d08ded8da..668c5f0bb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,4 +1,7 @@ # Changelog +## 4.32 + - [#331] (https://github.com/cohere-ai/cohere-python/pull/331) + - Embed: add `compression` parameter for embed models ## 4.31 - [#324] (https://github.com/cohere-ai/cohere-python/pull/324) diff --git a/pyproject.toml b/pyproject.toml index dfbec6d7a..d00921663 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "cohere" -version = "4.31" +version = "4.32" description = "" authors = ["Cohere"] readme = "README.md" From 4ae2f3fb4650da2efe74189bebb95e28bdbae8d1 Mon Sep 17 00:00:00 2001 From: Alekhya Date: Mon, 23 Oct 2023 20:28:23 -0700 Subject: [PATCH 3/6] remove compress codebook from embed job --- cohere/client.py | 6 ------ cohere/client_async.py | 6 ------ 2 files changed, 12 deletions(-) diff --git a/cohere/client.py b/cohere/client.py index 5d671381b..7878ae6ea 100644 --- a/cohere/client.py +++ b/cohere/client.py @@ -1044,8 +1044,6 @@ def create_embed_job( name: Optional[str] = None, model: Optional[str] = None, truncate: Optional[str] = None, - compress: Optional[bool] = None, - compression_codebook: Optional[str] = None, text_field: Optional[str] = None, ) -> EmbedJob: """Create embed job. @@ -1055,8 +1053,6 @@ def create_embed_job( name (Optional[str], optional): The name of the embed job. Defaults to None. model (Optional[str], optional): The model ID to use for embedding the text. Defaults to None. truncate (Optional[str], optional): How the API handles text longer than the maximum token length. Defaults to None. - compress (Optional[bool], optional): Use embedding compression. Defaults to None. - compression_codebook (Optional[str], optional): Embedding compression codebook. Defaults to None. text_field (Optional[str], optional): Name of the column containing text to embed. Defaults to None. Returns: @@ -1075,8 +1071,6 @@ def create_embed_job( "name": name, "model": model, "truncate": truncate, - "compress": compress, - "compression_codebook": compression_codebook, "text_field": text_field, "output_format": "avro", } diff --git a/cohere/client_async.py b/cohere/client_async.py index f67708caa..29c363d5b 100644 --- a/cohere/client_async.py +++ b/cohere/client_async.py @@ -722,8 +722,6 @@ async def create_embed_job( name: Optional[str] = None, model: Optional[str] = None, truncate: Optional[str] = None, - compress: Optional[bool] = None, - compression_codebook: Optional[str] = None, text_field: Optional[str] = None, ) -> AsyncEmbedJob: """Create embed job. @@ -733,8 +731,6 @@ async def create_embed_job( name (Optional[str], optional): The name of the embed job. Defaults to None. model (Optional[str], optional): The model ID to use for embedding the text. Defaults to None. truncate (Optional[str], optional): How the API handles text longer than the maximum token length. Defaults to None. - compress (Optional[bool], optional): Use embedding compression. Defaults to None. - compression_codebook (Optional[str], optional): Embedding compression codebook. Defaults to None. text_field (Optional[str], optional): Name of the column containing text to embed. Defaults to None. Returns: @@ -753,8 +749,6 @@ async def create_embed_job( "name": name, "model": model, "truncate": truncate, - "compress": compress, - "compression_codebook": compression_codebook, "text_field": text_field, "output_format": "avro", } From ab3ce5fef49075922b0b9d790034aa2c295dee7b Mon Sep 17 00:00:00 2001 From: Alekhya Date: Mon, 23 Oct 2023 20:31:15 -0700 Subject: [PATCH 4/6] remove test for compress codebook --- tests/async/test_async_codebook.py | 37 ---------------------------- tests/sync/test_codebook.py | 39 ------------------------------ 2 files changed, 76 deletions(-) delete mode 100644 tests/async/test_async_codebook.py delete mode 100644 tests/sync/test_codebook.py diff --git a/tests/async/test_async_codebook.py b/tests/async/test_async_codebook.py deleted file mode 100644 index 04f468014..000000000 --- a/tests/async/test_async_codebook.py +++ /dev/null @@ -1,37 +0,0 @@ -import pytest - -from cohere.responses import Codebook - - -@pytest.mark.asyncio -async def test_async_multiple_codebook(async_client): - compression_codebooks = { - "PQ96": (96, 256, 8), - "PQ64": (64, 256, 12), - "PQ48": (48, 256, 16), - "PQ32": (32, 256, 24), - } - for cb, (segments, num_centroids, length) in compression_codebooks.items(): - codebook = await async_client.codebook( - model="multilingual-22-12", - compression_codebook=cb, - ) - assert isinstance(codebook, Codebook) - assert isinstance(codebook.codebook[0], list) - assert isinstance(codebook.codebook[0][0], list) - assert len(codebook.codebook) == segments - assert len(codebook.codebook[0]) == num_centroids - assert len(codebook.codebook[0][0]) == length - assert codebook.meta - assert codebook.meta["api_version"]["version"] - assert codebook.meta["api_version"] - - -@pytest.mark.asyncio -async def test_async_codebook_default(async_client): - codebook = await async_client.codebook(model="multilingual-22-12") - assert len(codebook) == 96 - assert isinstance(codebook, Codebook) - assert codebook.meta - assert codebook.meta["api_version"] - assert codebook.meta["api_version"]["version"] diff --git a/tests/sync/test_codebook.py b/tests/sync/test_codebook.py deleted file mode 100644 index 8c9fd7a14..000000000 --- a/tests/sync/test_codebook.py +++ /dev/null @@ -1,39 +0,0 @@ -import unittest - -from utils import get_api_key - -import cohere - -API_KEY = get_api_key() -co = cohere.Client(API_KEY) - - -class TestCodebook(unittest.TestCase): - def test_success(self): - compression_codebooks = { - "PQ96": (96, 256, 8), - "PQ64": (64, 256, 12), - "PQ48": (48, 256, 16), - "PQ32": (32, 256, 24), - } - for cb, (segments, num_centroids, length) in compression_codebooks.items(): - prediction = co.codebook( - model="multilingual-22-12", - compression_codebook=cb, - ) - self.assertIsInstance(prediction.codebook[0], list) - self.assertIsInstance(prediction.codebook[0][0], list) - self.assertEqual(len(prediction.codebook), segments) - self.assertEqual(len(prediction.codebook[0]), num_centroids) - self.assertEqual(len(prediction.codebook[0][0]), length) - self.assertTrue(prediction.meta) - self.assertTrue(prediction.meta["api_version"]) - self.assertTrue(prediction.meta["api_version"]["version"]) - - def test_default_model(self): - prediction = co.codebook( - model="multilingual-22-12", - ) - self.assertEqual(len(prediction.codebook), 96) - self.assertIsInstance(prediction.codebook[0], list) - self.assertIsInstance(prediction.codebook[0][0], list) From 8b98e82b05974dd0626a1d1c502a6017ad61c1a6 Mon Sep 17 00:00:00 2001 From: Alekhya Date: Mon, 23 Oct 2023 20:39:13 -0700 Subject: [PATCH 5/6] fix test --- cohere/client_async.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cohere/client_async.py b/cohere/client_async.py index 29c363d5b..bb192fde6 100644 --- a/cohere/client_async.py +++ b/cohere/client_async.py @@ -298,7 +298,7 @@ async def embed( return Embeddings( embeddings=[e for res in responses for e in res["embeddings"]], - compressed_embeddings=[e for res in responses for e in res["compressed_embeddings"]] if compress else None, + compressed_embeddings=[e for res in responses for e in res["compressed_embeddings"]] if compression else None, meta=meta, ) From 1d648c9a21f37cae9ab00e4107b32644c08800cc Mon Sep 17 00:00:00 2001 From: Alekhya Date: Mon, 23 Oct 2023 20:49:57 -0700 Subject: [PATCH 6/6] lint --- cohere/client_async.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/cohere/client_async.py b/cohere/client_async.py index bb192fde6..132fd73a5 100644 --- a/cohere/client_async.py +++ b/cohere/client_async.py @@ -298,7 +298,9 @@ async def embed( return Embeddings( embeddings=[e for res in responses for e in res["embeddings"]], - compressed_embeddings=[e for res in responses for e in res["compressed_embeddings"]] if compression else None, + compressed_embeddings=[e for res in responses for e in res["compressed_embeddings"]] + if compression + else None, meta=meta, )