From 0f73ba96776085502deed4d41a72633428d9dfc1 Mon Sep 17 00:00:00 2001 From: Sukkrit Sharma <44338419+sukkritsharmaofficial@users.noreply.github.com> Date: Fri, 22 Dec 2023 05:10:00 +0530 Subject: [PATCH] Added Support for Ollama for local model inference. (#1045) Co-authored-by: Deshraj Yadav --- configs/ollama.yaml | 12 ++ docs/components/data-sources/docs-site.mdx | 2 +- docs/components/llms.mdx | 27 +++ docs/examples/notebooks-and-replits.mdx | 4 + embedchain/factory.py | 1 + embedchain/llm/ollama.py | 34 +++ embedchain/models/providers.py | 1 + embedchain/utils.py | 1 + notebooks/ollama.ipynb | 228 +++++++++++++++++++++ tests/llm/test_ollama.py | 38 ++++ tests/test_utils.py | 1 + 11 files changed, 348 insertions(+), 1 deletion(-) create mode 100644 configs/ollama.yaml create mode 100644 embedchain/llm/ollama.py create mode 100644 notebooks/ollama.ipynb create mode 100644 tests/llm/test_ollama.py diff --git a/configs/ollama.yaml b/configs/ollama.yaml new file mode 100644 index 00000000..5b01cc18 --- /dev/null +++ b/configs/ollama.yaml @@ -0,0 +1,12 @@ +llm: + provider: ollama + config: + model: 'llama2' + temperature: 0.5 + top_p: 1 + stream: true + +embedder: + provider: huggingface + config: + model: 'BAAI/bge-small-en-v1.5' diff --git a/docs/components/data-sources/docs-site.mdx b/docs/components/data-sources/docs-site.mdx index 0d3fcd9e..14e4a00d 100644 --- a/docs/components/data-sources/docs-site.mdx +++ b/docs/components/data-sources/docs-site.mdx @@ -10,5 +10,5 @@ from embedchain import Pipeline as App app = App() app.add("https://docs.embedchain.ai/", data_type="docs_site") app.query("What is Embedchain?") -# Answer: Embedchain is a platform that utilizes various components, including paid/proprietary ones, to provide what is believed to be the best configuration available. It uses LLM (Language Model) providers such as OpenAI, Anthpropic, Vertex_AI, GPT4ALL, Azure_OpenAI, LLAMA2, JINA, and COHERE. Embedchain allows users to import and utilize these LLM providers for their applications.' +# Answer: Embedchain is a platform that utilizes various components, including paid/proprietary ones, to provide what is believed to be the best configuration available. It uses LLM (Language Model) providers such as OpenAI, Anthpropic, Vertex_AI, GPT4ALL, Azure_OpenAI, LLAMA2, JINA, Ollama and COHERE. Embedchain allows users to import and utilize these LLM providers for their applications.' ``` diff --git a/docs/components/llms.mdx b/docs/components/llms.mdx index a48145c5..7cc43fa3 100644 --- a/docs/components/llms.mdx +++ b/docs/components/llms.mdx @@ -12,6 +12,7 @@ Embedchain comes with built-in support for various popular large language models + @@ -329,6 +330,32 @@ llm: +## Ollama + +Setup Ollama using https://github.com/jmorganca/ollama + + + +```python main.py +import os +from embedchain import Pipeline as App + +# load llm configuration from config.yaml file +app = App.from_config(config_path="config.yaml") +``` + +```yaml config.yaml +llm: + provider: ollama + config: + model: 'llama2' + temperature: 0.5 + top_p: 1 + stream: true +``` + + + ## GPT4ALL Install related dependencies using the following command: diff --git a/docs/examples/notebooks-and-replits.mdx b/docs/examples/notebooks-and-replits.mdx index 82da96fb..ef09c21d 100644 --- a/docs/examples/notebooks-and-replits.mdx +++ b/docs/examples/notebooks-and-replits.mdx @@ -44,6 +44,10 @@ Get started with Embedchain by trying out the examples below. You can run the ex Open In Colab Try with Replit Badge + + Ollama + Open In Colab + Hugging Face Open In Colab diff --git a/embedchain/factory.py b/embedchain/factory.py index 6ec1479b..b77d71ab 100644 --- a/embedchain/factory.py +++ b/embedchain/factory.py @@ -13,6 +13,7 @@ class LlmFactory: "azure_openai": "embedchain.llm.azure_openai.AzureOpenAILlm", "cohere": "embedchain.llm.cohere.CohereLlm", "gpt4all": "embedchain.llm.gpt4all.GPT4ALLLlm", + "ollama": "embedchain.llm.ollama.OllamaLlm", "huggingface": "embedchain.llm.huggingface.HuggingFaceLlm", "jina": "embedchain.llm.jina.JinaLlm", "llama2": "embedchain.llm.llama2.Llama2Llm", diff --git a/embedchain/llm/ollama.py b/embedchain/llm/ollama.py new file mode 100644 index 00000000..459a22b6 --- /dev/null +++ b/embedchain/llm/ollama.py @@ -0,0 +1,34 @@ +from typing import Iterable, Optional, Union + +from langchain.callbacks.manager import CallbackManager +from langchain.callbacks.stdout import StdOutCallbackHandler +from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler +from langchain.llms.ollama import Ollama + +from embedchain.config import BaseLlmConfig +from embedchain.helpers.json_serializable import register_deserializable +from embedchain.llm.base import BaseLlm + + +@register_deserializable +class OllamaLlm(BaseLlm): + def __init__(self, config: Optional[BaseLlmConfig] = None): + super().__init__(config=config) + if self.config.model is None: + self.config.model = "llama2" + + def get_llm_model_answer(self, prompt): + return self._get_answer(prompt=prompt, config=self.config) + + def _get_answer(self, prompt: str, config: BaseLlmConfig) -> Union[str, Iterable]: + callback_manager = [StreamingStdOutCallbackHandler()] if config.stream else [StdOutCallbackHandler()] + + llm = Ollama( + model=config.model, + system=config.system_prompt, + temperature=config.temperature, + top_p=config.top_p, + callback_manager=CallbackManager(callback_manager) + ) + + return llm(prompt) diff --git a/embedchain/models/providers.py b/embedchain/models/providers.py index bd019cff..62c93675 100644 --- a/embedchain/models/providers.py +++ b/embedchain/models/providers.py @@ -6,4 +6,5 @@ class Providers(Enum): ANTHROPHIC = "ANTHPROPIC" VERTEX_AI = "VERTEX_AI" GPT4ALL = "GPT4ALL" + OLLAMA = "OLLAMA" AZURE_OPENAI = "AZURE_OPENAI" diff --git a/embedchain/utils.py b/embedchain/utils.py index 3838d805..8885f76e 100644 --- a/embedchain/utils.py +++ b/embedchain/utils.py @@ -385,6 +385,7 @@ def validate_config(config_data): "huggingface", "cohere", "gpt4all", + "ollama", "jina", "llama2", "vertexai", diff --git a/notebooks/ollama.ipynb b/notebooks/ollama.ipynb new file mode 100644 index 00000000..8abd12d4 --- /dev/null +++ b/notebooks/ollama.ipynb @@ -0,0 +1,228 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "id": "b02n_zJ_hl3d" + }, + "source": [ + "## Cookbook for using Ollama with Embedchain" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "gyJ6ui2vhtMY" + }, + "source": [ + "### Step-1: Setup Ollama, follow these instructions https://github.com/jmorganca/ollama\n", + "\n", + "Once Setup is done:\n", + "\n", + "- ollama pull llama2 (All supported models can be found here: https://ollama.ai/library)\n", + "- ollama run llama2 (Test out the model once)\n", + "- ollama serve" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "Ns6RhPfbiitr" + }, + "source": [ + "### Step-2: Define your llm and embedding model config (Going all out local inference, no need for OpenAI API Key)" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": { + "id": "S9CkxVjriotB" + }, + "outputs": [], + "source": [ + "config = \"\"\"\n", + "llm:\n", + " provider: ollama\n", + " config:\n", + " model: 'llama2'\n", + " temperature: 0.5\n", + " top_p: 1\n", + " stream: true\n", + "\n", + "embedder:\n", + " provider: huggingface\n", + " config:\n", + " model: 'BAAI/bge-small-en-v1.5'\n", + "\"\"\"\n", + "\n", + "# Write the multi-line string to a YAML file\n", + "with open('ollama.yaml', 'w') as file:\n", + " file.write(config)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "PGt6uPLIi1CS" + }, + "source": [ + "### Step-3 Create embedchain app based on the config" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/", + "height": 321 + }, + "id": "Amzxk3m-i3tD", + "outputId": "afe8afde-5cb8-46bc-c541-3ad26cc3fa6e" + }, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/Users/sukkritsharma/workspace/embedchain/.venv/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n", + " from .autonotebook import tqdm as notebook_tqdm\n" + ] + } + ], + "source": [ + "from embedchain import Pipeline as App\n", + "app = App.from_config(config_path=\"ollama.yaml\")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "XNXv4yZwi7ef" + }, + "source": [ + "### Step-4: Add data sources to your app" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/", + "height": 176 + }, + "id": "Sn_0rx9QjIY9", + "outputId": "2f2718a4-3b7e-4844-fd46-3e0857653ca0" + }, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Inserting batches in chromadb: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00, 1.57it/s]" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Successfully saved https://www.forbes.com/profile/elon-musk (DataType.WEB_PAGE). New chunks count: 4\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "\n" + ] + }, + { + "data": { + "text/plain": [ + "'8cf46026cabf9b05394a2658bd1fe890'" + ] + }, + "execution_count": 3, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "app.add(\"https://www.forbes.com/profile/elon-musk\")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "_7W6fDeAjMAP" + }, + "source": [ + "### Step-5: All set. Now start asking questions related to your data" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "cvIK7dWRjN_f", + "outputId": "79e873c8-9594-45da-f5a3-0a893511267f" + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Elon Musk is a business magnate, investor, and engineer. He is the CEO of SpaceX and Tesla, Inc., and has been involved in other successful ventures such as Neuralink and The Boring Company. Musk is known for his innovative ideas, entrepreneurial spirit, and vision for the future of humanity.\n", + "\n", + "As the CEO of Tesla, Musk has played a significant role in popularizing electric vehicles and making them more accessible to the masses. Under his leadership, Tesla has grown into one of the most valuable companies in the world.\n", + "\n", + "SpaceX, another company founded by Musk, is a leading player in the commercial space industry. SpaceX has developed advanced rockets and spacecraft, including the Falcon 9 and Dragon, which have successfully launched numerous satellites and other payloads into orbit.\n", + "\n", + "Musk is also known for his ambitious goals, such as establishing a human settlement on Mars and developing sustainable energy solutions to address climate change. He has been recognized for his philanthropic efforts, particularly in the area of education, and has been awarded numerous honors and awards for his contributions to society.\n", + "\n", + "Overall, Elon Musk is a highly influential and innovative entrepreneur who has made significant impacts in various industries and has inspired many people around the world with his vision and leadership." + ] + } + ], + "source": [ + "answer = app.query(\"who is elon musk?\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "colab": { + "provenance": [] + }, + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.9" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/tests/llm/test_ollama.py b/tests/llm/test_ollama.py new file mode 100644 index 00000000..394bbd41 --- /dev/null +++ b/tests/llm/test_ollama.py @@ -0,0 +1,38 @@ +import pytest + +from embedchain.config import BaseLlmConfig +from embedchain.llm.ollama import OllamaLlm + + +@pytest.fixture +def ollama_llm_config(): + config = BaseLlmConfig(model="llama2", temperature=0.7, top_p=0.8, stream=True, system_prompt=None) + yield config + +def test_get_llm_model_answer(ollama_llm_config, mocker): + mocker.patch("embedchain.llm.ollama.OllamaLlm._get_answer", return_value="Test answer") + + llm = OllamaLlm(ollama_llm_config) + answer = llm.get_llm_model_answer("Test query") + + assert answer == "Test answer" + + +def test_get_answer_mocked_ollama(ollama_llm_config, mocker): + mocked_ollama = mocker.patch("embedchain.llm.ollama.Ollama") + mock_instance = mocked_ollama.return_value + mock_instance.return_value = "Mocked answer" + + llm = OllamaLlm(ollama_llm_config) + prompt = "Test query" + answer = llm.get_llm_model_answer(prompt) + + assert answer == "Mocked answer" + mocked_ollama.assert_called_once_with( + model="llama2", + system=None, + temperature=0.7, + top_p=0.8, + callback_manager=mocker.ANY # Use mocker.ANY to ignore the exact instance + ) + mock_instance.assert_called_once_with(prompt) diff --git a/tests/test_utils.py b/tests/test_utils.py index 891bb453..92ab6bd1 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -8,6 +8,7 @@ CONFIG_YAMLS = [ "configs/chroma.yaml", "configs/chunker.yaml", "configs/cohere.yaml", + "configs/ollama.yaml", "configs/full-stack.yaml", "configs/gpt4.yaml", "configs/gpt4all.yaml",