[embed] add compression parameter (#331)

* add compression parameter to embed
* remove compress codebook
* update changelog and toml

Author: alekhya-n

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)

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,
                 }
             )
@@ -1047,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.
@@ -1058,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:
@@ -1078,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",
         }

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,17 +280,15 @@ 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 = [
             dict(
                 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)
@@ -301,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 compress else None,
+            compressed_embeddings=[e for res in responses for e in res["compressed_embeddings"]]
+            if compression
+            else None,
             meta=meta,
         )
 
@@ -725,8 +724,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.
@@ -736,8 +733,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:
@@ -756,8 +751,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",
         }

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "cohere"
-version = "4.31"
+version = "4.32"
 description = ""
 authors = ["Cohere"]
 readme = "README.md"