From 804e9d2f2c75c1a1ebba0283e08134c95e378462 Mon Sep 17 00:00:00 2001
From: Piyush Jain <piyushjain@duck.com>
Date: Wed, 8 Nov 2023 13:44:38 -0800
Subject: [PATCH] Added models for gpt4all.

---
 .../jupyter-ai-magics/jupyter_ai_magics/__init__.py  |  1 +
 .../jupyter_ai_magics/embedding_providers.py         |  9 +++++++++
 .../jupyter-ai-magics/jupyter_ai_magics/providers.py | 12 ++++++++++++
 packages/jupyter-ai-magics/pyproject.toml            |  1 +
 4 files changed, 23 insertions(+)

diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/__init__.py b/packages/jupyter-ai-magics/jupyter_ai_magics/__init__.py
index ba756a452..b991a49b9 100644
--- a/packages/jupyter-ai-magics/jupyter_ai_magics/__init__.py
+++ b/packages/jupyter-ai-magics/jupyter_ai_magics/__init__.py
@@ -4,6 +4,7 @@
 from .embedding_providers import (
     BedrockEmbeddingsProvider,
     CohereEmbeddingsProvider,
+    GPT4AllEmbeddingsProvider,
     HfHubEmbeddingsProvider,
     OpenAIEmbeddingsProvider,
 )
diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/embedding_providers.py b/packages/jupyter-ai-magics/jupyter_ai_magics/embedding_providers.py
index d09907fb0..1fba25f50 100644
--- a/packages/jupyter-ai-magics/jupyter_ai_magics/embedding_providers.py
+++ b/packages/jupyter-ai-magics/jupyter_ai_magics/embedding_providers.py
@@ -9,6 +9,7 @@
 from langchain.embeddings import (
     BedrockEmbeddings,
     CohereEmbeddings,
+    GPT4AllEmbeddings,
     HuggingFaceHubEmbeddings,
     OpenAIEmbeddings,
 )
@@ -103,3 +104,11 @@ class BedrockEmbeddingsProvider(BaseEmbeddingsProvider, BedrockEmbeddings):
     model_id_key = "model_id"
     pypi_package_deps = ["boto3"]
     auth_strategy = AwsAuthStrategy()
+
+
+class GPT4AllEmbeddingsProvider(BaseEmbeddingsProvider, GPT4AllEmbeddings):
+    id = "gpt4all"
+    name = "GPT4All Embeddings"
+    models = ["all-MiniLM-L6-v2-f16"]
+    model_id_key = "model_id"
+    pypi_package_deps = ["gpt4all"]
\ No newline at end of file
diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py b/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py
index a17067921..e997a4312 100644
--- a/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py
+++ b/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py
@@ -342,6 +342,18 @@ def __init__(self, **kwargs):
         "ggml-gpt4all-j-v1.3-groovy",
         # this one needs llama backend and has licence restriction
         "ggml-gpt4all-l13b-snoozy",
+        "mistral-7b-openorca.Q4_0",
+        "mistral-7b-instruct-v0.1.Q4_0",
+        "gpt4all-falcon-q4_0",
+        "wizardlm-13b-v1.2.Q4_0",
+        "nous-hermes-llama2-13b.Q4_0",
+        "gpt4all-13b-snoozy-q4_0",
+        "mpt-7b-chat-merges-q4_0",
+        "orca-mini-3b-gguf2-q4_0",
+        "replit-code-v1_5-3b-q4_0",
+        "starcoder-q4_0",
+        "rift-coder-v0-7b-q4_0",
+        "em_german_mistral_v01.Q4_0"
     ]
     model_id_key = "model"
     pypi_package_deps = ["gpt4all"]
diff --git a/packages/jupyter-ai-magics/pyproject.toml b/packages/jupyter-ai-magics/pyproject.toml
index 7eb8a8644..60e231047 100644
--- a/packages/jupyter-ai-magics/pyproject.toml
+++ b/packages/jupyter-ai-magics/pyproject.toml
@@ -72,6 +72,7 @@ amazon-bedrock-chat = "jupyter_ai_magics:BedrockChatProvider"
 [project.entry-points."jupyter_ai.embeddings_model_providers"]
 bedrock = "jupyter_ai_magics:BedrockEmbeddingsProvider"
 cohere = "jupyter_ai_magics:CohereEmbeddingsProvider"
+gpt4all = "jupyter_ai_magics:GPT4AllEmbeddingsProvider"
 huggingface_hub = "jupyter_ai_magics:HfHubEmbeddingsProvider"
 openai = "jupyter_ai_magics:OpenAIEmbeddingsProvider"