Upgrades cohere dependency, model list (#594)

jupyterlab · Jan 22, 2024 · b2e7c95 · b2e7c95
1 parent f4de0c8
commit b2e7c95
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 2 deletions.
diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/embedding_providers.py b/packages/jupyter-ai-magics/jupyter_ai_magics/embedding_providers.py
@@ -80,7 +80,15 @@ class OpenAIEmbeddingsProvider(BaseEmbeddingsProvider, OpenAIEmbeddings):
 class CohereEmbeddingsProvider(BaseEmbeddingsProvider, CohereEmbeddings):
     id = "cohere"
     name = "Cohere"
-    models = ["large", "multilingual-22-12", "small"]
+    models = [
+        "embed-english-v2.0",
+        "embed-english-light-v2.0",
+        "embed-multilingual-v2.0",
+        "embed-english-v3.0",
+        "embed-english-light-v3.0",
+        "embed-multilingual-v3.0",
+        "embed-multilingual-light-v3.0",
+    ]
     model_id_key = "model"
     pypi_package_deps = ["cohere"]
     auth_strategy = EnvAuthStrategy(name="COHERE_API_KEY")

diff --git a/packages/jupyter-ai-magics/pyproject.toml b/packages/jupyter-ai-magics/pyproject.toml
@@ -37,7 +37,7 @@ test = ["coverage", "pytest", "pytest-asyncio", "pytest-cov"]
 all = [
     "ai21",
     "anthropic~=0.3.0",
-    "cohere",
+    "cohere>4.40,<5",
     "gpt4all",
     "huggingface_hub",
     "ipywidgets",