From 9f70ea7de675f136f544419e13249748ad64fb35 Mon Sep 17 00:00:00 2001 From: Jeremy Tuloup Date: Thu, 11 Jul 2024 00:06:17 +0200 Subject: [PATCH] Add Ollama (#646) * Add Ollama * mistral:text for embeddings * Add gemma * Update the list of models Co-authored-by: Bc * Mention Ollama in the docs * Apply suggestions from code review Co-authored-by: Piyush Jain Co-authored-by: david qiu --------- Co-authored-by: Bc Co-authored-by: Piyush Jain Co-authored-by: david qiu --- README.md | 2 +- docs/source/users/index.md | 4 ++++ .../jupyter_ai_magics/__init__.py | 2 ++ .../jupyter_ai_magics/embedding_providers.py | 14 ++++++++++++++ .../jupyter_ai_magics/providers.py | 13 +++++++++++++ packages/jupyter-ai-magics/pyproject.toml | 2 ++ 6 files changed, 36 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 41be45c6b..682c4bf23 100644 --- a/README.md +++ b/README.md @@ -11,7 +11,7 @@ in JupyterLab and the Jupyter Notebook. More specifically, Jupyter AI offers: * A native chat UI in JupyterLab that enables you to work with generative AI as a conversational assistant. * Support for a wide range of generative model providers, including AI21, Anthropic, AWS, Cohere, Gemini, Hugging Face, MistralAI, NVIDIA, and OpenAI. -* Local model support through GPT4All, enabling use of generative AI models on consumer grade machines +* Local model support through GPT4All and Ollama, enabling use of generative AI models on consumer grade machines with ease and privacy. Documentation is available on [ReadTheDocs](https://jupyter-ai.readthedocs.io/en/latest/). diff --git a/docs/source/users/index.md b/docs/source/users/index.md index 8db821238..2365d4784 100644 --- a/docs/source/users/index.md +++ b/docs/source/users/index.md @@ -357,6 +357,10 @@ GPT4All support is still an early-stage feature, so some bugs may be encountered during usage. Our team is still actively improving support for locally-hosted models. +### Ollama usage + +To get started, follow the instructions on the [Ollama website](https://ollama.com/) to set up `ollama` and download the models locally. To select a model, enter the model name in the settings panel, for example `deepseek-coder-v2`. + ### Asking about something in your notebook Jupyter AI's chat interface can include a portion of your notebook in your prompt. diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/__init__.py b/packages/jupyter-ai-magics/jupyter_ai_magics/__init__.py index d6ca65135..f596c08cd 100644 --- a/packages/jupyter-ai-magics/jupyter_ai_magics/__init__.py +++ b/packages/jupyter-ai-magics/jupyter_ai_magics/__init__.py @@ -6,6 +6,7 @@ BedrockEmbeddingsProvider, GPT4AllEmbeddingsProvider, HfHubEmbeddingsProvider, + OllamaEmbeddingsProvider, QianfanEmbeddingsEndpointProvider, ) from .exception import store_exception @@ -23,6 +24,7 @@ BedrockProvider, GPT4AllProvider, HfHubProvider, + OllamaProvider, QianfanProvider, SmEndpointProvider, TogetherAIProvider, 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 9f3bde6e9..6a175193b 100644 --- a/packages/jupyter-ai-magics/jupyter_ai_magics/embedding_providers.py +++ b/packages/jupyter-ai-magics/jupyter_ai_magics/embedding_providers.py @@ -12,6 +12,7 @@ BedrockEmbeddings, GPT4AllEmbeddings, HuggingFaceHubEmbeddings, + OllamaEmbeddings, QianfanEmbeddingsEndpoint, ) @@ -66,6 +67,19 @@ def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs, **model_kwargs) +class OllamaEmbeddingsProvider(BaseEmbeddingsProvider, OllamaEmbeddings): + id = "ollama" + name = "Ollama" + # source: https://ollama.com/library + models = [ + "nomic-embed-text", + "mxbai-embed-large", + "all-minilm", + "snowflake-arctic-embed", + ] + model_id_key = "model" + + class HfHubEmbeddingsProvider(BaseEmbeddingsProvider, HuggingFaceHubEmbeddings): id = "huggingface_hub" name = "Hugging Face Hub" diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py b/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py index 39b1b8adf..e017125ba 100644 --- a/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py +++ b/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py @@ -37,6 +37,7 @@ Bedrock, GPT4All, HuggingFaceEndpoint, + Ollama, SagemakerEndpoint, Together, ) @@ -723,6 +724,18 @@ async def _acall(self, *args, **kwargs) -> Coroutine[Any, Any, str]: return await self._call_in_executor(*args, **kwargs) +class OllamaProvider(BaseProvider, Ollama): + id = "ollama" + name = "Ollama" + model_id_key = "model" + help = ( + "See [https://www.ollama.com/library](https://www.ollama.com/library) for a list of models. " + "Pass a model's name; for example, `deepseek-coder-v2`." + ) + models = ["*"] + registry = True + + class JsonContentHandler(LLMContentHandler): content_type = "application/json" accepts = "application/json" diff --git a/packages/jupyter-ai-magics/pyproject.toml b/packages/jupyter-ai-magics/pyproject.toml index 5e3791402..3deff9228 100644 --- a/packages/jupyter-ai-magics/pyproject.toml +++ b/packages/jupyter-ai-magics/pyproject.toml @@ -58,6 +58,7 @@ anthropic-chat = "jupyter_ai_magics.partner_providers.anthropic:ChatAnthropicPro cohere = "jupyter_ai_magics.partner_providers.cohere:CohereProvider" gpt4all = "jupyter_ai_magics:GPT4AllProvider" huggingface_hub = "jupyter_ai_magics:HfHubProvider" +ollama = "jupyter_ai_magics:OllamaProvider" openai = "jupyter_ai_magics.partner_providers.openai:OpenAIProvider" openai-chat = "jupyter_ai_magics.partner_providers.openai:ChatOpenAIProvider" azure-chat-openai = "jupyter_ai_magics.partner_providers.openai:AzureChatOpenAIProvider" @@ -76,6 +77,7 @@ cohere = "jupyter_ai_magics.partner_providers.cohere:CohereEmbeddingsProvider" mistralai = "jupyter_ai_magics.partner_providers.mistralai:MistralAIEmbeddingsProvider" gpt4all = "jupyter_ai_magics:GPT4AllEmbeddingsProvider" huggingface_hub = "jupyter_ai_magics:HfHubEmbeddingsProvider" +ollama = "jupyter_ai_magics:OllamaEmbeddingsProvider" openai = "jupyter_ai_magics.partner_providers.openai:OpenAIEmbeddingsProvider" qianfan = "jupyter_ai_magics:QianfanEmbeddingsEndpointProvider"