Remove old Cohere models and add aliases for existing ones (#6007)

* Remove old cohere models * Add aliases for the existing models according to Cohere documentation * Add release note * put cohere embdding models in a constant * update doc strings --------- Co-authored-by: Massimiliano Pippi <[email protected]>
deepset-ai · Oct 13, 2023 · ad25041 · ad25041
1 parent fbd22bc
commit ad25041
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 10 deletions.
diff --git a/haystack/nodes/retriever/_embedding_encoder.py b/haystack/nodes/retriever/_embedding_encoder.py
@@ -46,6 +46,14 @@
 COHERE_TIMEOUT = float(os.environ.get(HAYSTACK_REMOTE_API_TIMEOUT_SEC, 30))
 COHERE_BACKOFF = int(os.environ.get(HAYSTACK_REMOTE_API_BACKOFF_SEC, 10))
 COHERE_MAX_RETRIES = int(os.environ.get(HAYSTACK_REMOTE_API_MAX_RETRIES, 5))
+COHERE_EMBEDDING_MODELS = [
+    "small",
+    "large",
+    "multilingual-22-12",
+    "embed-english-v2.0",
+    "embed-english-light-v2.0",
+    "embed-multilingual-v2.0",
+]
 
 
 class _DefaultEmbeddingEncoder(_BaseEmbeddingEncoder):
@@ -364,20 +372,15 @@ class _CohereEmbeddingEncoder(_BaseEmbeddingEncoder):
     def __init__(self, retriever: "EmbeddingRetriever"):
         torch_and_transformers_import.check()
 
-        # See https://docs.cohere.ai/embed-reference/ for more details
+        # See https://docs.cohere.com/reference/embed for more details
         # Cohere has a max seq length of 4096 tokens and a max batch size of 96
         self.max_seq_len = min(4096, retriever.max_seq_len)
         self.url = "https://api.cohere.ai/embed"
         self.api_key = retriever.api_key
         self.batch_size = min(96, retriever.batch_size)
         self.progress_bar = retriever.progress_bar
         self.model: str = next(
-            (
-                m
-                for m in ["small", "medium", "large", "multilingual-22-12", "finance-sentiment"]
-                if m in retriever.embedding_model
-            ),
-            "multilingual-22-12",
+            (m for m in COHERE_EMBEDDING_MODELS if m in retriever.embedding_model), "multilingual-22-12"
         )
 
     @retry(

diff --git a/haystack/nodes/retriever/dense.py b/haystack/nodes/retriever/dense.py
@@ -17,7 +17,7 @@
 from haystack.schema import Document, FilterType
 from haystack.document_stores import BaseDocumentStore
 from haystack.nodes.retriever.base import BaseRetriever
-from haystack.nodes.retriever._embedding_encoder import _EMBEDDING_ENCODERS
+from haystack.nodes.retriever._embedding_encoder import _EMBEDDING_ENCODERS, COHERE_EMBEDDING_MODELS
 from haystack.utils.early_stopping import EarlyStopping
 from haystack.telemetry import send_event
 from haystack.lazy_imports import LazyImport
@@ -1472,7 +1472,7 @@ def __init__(
         :param embedding_model: Local path or name of model in Hugging Face's model hub such
                                 as ``'sentence-transformers/all-MiniLM-L6-v2'``. The embedding model could also
                                 potentially be an OpenAI model ["ada", "babbage", "davinci", "curie"] or
-                                a Cohere model ["small", "medium", "large"].
+                                a Cohere model ["embed-english-v2.0", "embed-english-light-v2.0", "embed-multilingual-v2.0"].
         :param model_version: The version of model to use from the HuggingFace model hub. Can be tag name, branch name, or commit hash.
         :param use_gpu: Whether to use all available GPUs or the CPU. Falls back on CPU if no GPU is available.
         :param batch_size: Number of documents to encode at once.
@@ -1878,7 +1878,7 @@ def _infer_model_format(model_name_or_path: str, use_auth_token: Optional[Union[
         )
         if valid_openai_model_name:
             return "openai"
-        if model_name_or_path in ["small", "medium", "large", "multilingual-22-12", "finance-sentiment"]:
+        if model_name_or_path in COHERE_EMBEDDING_MODELS:
             return "cohere"
         # Check if model name is a local directory with sentence transformers config file in it
         if Path(model_name_or_path).exists():

diff --git a/releasenotes/notes/update-supported-cohere-models-68f17fa70051fc90.yaml b/releasenotes/notes/update-supported-cohere-models-68f17fa70051fc90.yaml
@@ -0,0 +1,7 @@
+---
+enhancements:
+  - |
+    Add alias names for Cohere embed models for an easier map between names
+fixes:
+  - |
+    Remove unsupported `medium` and `finance-sentiment` models from supported Cohere embed model list