Backport PR jupyterlab#594: Upgrades cohere dependency, model list (j…

…upyterlab#597) Co-authored-by: Jason Weill <[email protected]>
JasonWeill · Jan 23, 2024 · c8df21b · c8df21b
1 parent 5dec97e
commit c8df21b
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",