Update the list of models

Co-Authored-By: xiayu <[email protected]>
jupyterlab · Jul 10, 2024 · 697cf70 · 697cf70
1 parent 24a4500
commit 697cf70
Show file tree

Hide file tree

Showing 2 changed files with 16 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
@@ -71,9 +71,17 @@ class OllamaEmbeddingsProvider(BaseEmbeddingsProvider, OllamaEmbeddings):
     id = "ollama"
     name = "Ollama"
     models = [
+        "gemma",
+        "gemma2",
         "llama2",
-        "mistral:text",
+        "llama3",
+        "llama3:70b",
+        "phi3",
+        "mistral",
         "tinyllama",
+        "qwen2",
+        "qwen2:7b",
+        "qwen2:72b",
     ]
     model_id_key = "model"
 

diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py b/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py
@@ -728,12 +728,18 @@ class OllamaProvider(BaseProvider, Ollama):
     id = "ollama"
     name = "Ollama"
     model_id_key = "model"
-    # TODO: add more
     models = [
         "gemma",
+        "gemma2",
         "llama2",
+        "llama3",
+        "llama3:70b",
+        "phi3",
         "mistral",
         "tinyllama",
+        "qwen2",
+        "qwen2:7b",
+        "qwen2:72b",
     ]