From 4070fc1bf07e0025593ccf8d9c9205c7023ef5c1 Mon Sep 17 00:00:00 2001 From: Dev Khant Date: Sat, 8 Jun 2024 22:38:15 +0530 Subject: [PATCH] Fix ollama embeddings for remote machine (#1394) --- embedchain/config/embedder/ollama.py | 2 +- embedchain/embedder/ollama.py | 9 +++++---- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/embedchain/config/embedder/ollama.py b/embedchain/config/embedder/ollama.py index b08b39d7..c0a5f1de 100644 --- a/embedchain/config/embedder/ollama.py +++ b/embedchain/config/embedder/ollama.py @@ -12,4 +12,4 @@ class OllamaEmbedderConfig(BaseEmbedderConfig): base_url: Optional[str] = None, ): super().__init__(model) - self.base_url = base_url or "http://127.0.0.1:11434" + self.base_url = base_url or "http://localhost:11434" diff --git a/embedchain/embedder/ollama.py b/embedchain/embedder/ollama.py index a70e402e..9e4ada47 100644 --- a/embedchain/embedder/ollama.py +++ b/embedchain/embedder/ollama.py @@ -2,7 +2,7 @@ import logging from typing import Optional try: - import ollama + from ollama import Client except ImportError: raise ImportError("Ollama Embedder requires extra dependencies. Install with `pip install ollama`") from None @@ -19,11 +19,12 @@ class OllamaEmbedder(BaseEmbedder): def __init__(self, config: Optional[OllamaEmbedderConfig] = None): super().__init__(config=config) - local_models = ollama.list()["models"] + client = Client(host=config.base_url) + local_models = client.list()["models"] if not any(model.get("name") == self.config.model for model in local_models): logger.info(f"Pulling {self.config.model} from Ollama!") - ollama.pull(self.config.model) - embeddings = OllamaEmbeddings(model=self.config.model, base_url=self.config.base_url) + client.pull(self.config.model) + embeddings = OllamaEmbeddings(model=self.config.model, base_url=config.base_url) embedding_fn = BaseEmbedder._langchain_default_concept(embeddings) self.set_embedding_fn(embedding_fn=embedding_fn)