Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[embed] add compression parameter #331

Merged
merged 6 commits into from
Oct 24, 2023
Merged

[embed] add compression parameter #331

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
5791bc0
add compression parameter to embed
alekhya-n Oct 22, 2023
7298db5
update changelog and toml
alekhya-n Oct 22, 2023
4ae2f3f
remove compress codebook from embed job
alekhya-n Oct 24, 2023
ab3ce5f
remove test for compress codebook
alekhya-n Oct 24, 2023
8b98e82
fix test
alekhya-n Oct 24, 2023
1d648c9
lint
alekhya-n Oct 24, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions CHANGELOG.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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)
Expand Down
15 changes: 3 additions & 12 deletions cohere/client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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 = {
Expand All @@ -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,
}
)
Expand Down Expand Up @@ -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.
Expand All @@ -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:
Expand All @@ -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",
}
Expand Down
19 changes: 6 additions & 13 deletions cohere/client_async.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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)
Expand All @@ -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,
)

Expand Down Expand Up @@ -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.
Expand All @@ -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:
Expand All @@ -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",
}
Expand Down
2 changes: 1 addition & 1 deletion pyproject.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
[tool.poetry]
name = "cohere"
version = "4.31"
version = "4.32"
description = ""
authors = ["Cohere"]
readme = "README.md"
Expand Down
37 changes: 0 additions & 37 deletions tests/async/test_async_codebook.py
View file
Open in desktop

This file was deleted.

39 changes: 0 additions & 39 deletions tests/sync/test_codebook.py
View file
Open in desktop

This file was deleted.