From b43a116b3caf752afb2dd84519508bdd10412831 Mon Sep 17 00:00:00 2001 From: Dev Khant Date: Fri, 21 Jun 2024 21:26:46 +0530 Subject: [PATCH] Add vector dimension to Ollama embedder (#1435) --- embedchain/config/embedder/google.py | 3 ++- embedchain/config/embedder/ollama.py | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/embedchain/config/embedder/google.py b/embedchain/config/embedder/google.py index f42e5b53..7cf5a901 100644 --- a/embedchain/config/embedder/google.py +++ b/embedchain/config/embedder/google.py @@ -10,9 +10,10 @@ class GoogleAIEmbedderConfig(BaseEmbedderConfig): self, model: Optional[str] = None, deployment_name: Optional[str] = None, + vector_dimension: Optional[int] = None, task_type: Optional[str] = None, title: Optional[str] = None, ): - super().__init__(model, deployment_name) + super().__init__(model, deployment_name, vector_dimension) self.task_type = task_type or "retrieval_document" self.title = title or "Embeddings for Embedchain" diff --git a/embedchain/config/embedder/ollama.py b/embedchain/config/embedder/ollama.py index c0a5f1de..f680328f 100644 --- a/embedchain/config/embedder/ollama.py +++ b/embedchain/config/embedder/ollama.py @@ -10,6 +10,7 @@ class OllamaEmbedderConfig(BaseEmbedderConfig): self, model: Optional[str] = None, base_url: Optional[str] = None, + vector_dimension: Optional[int] = None, ): - super().__init__(model) + super().__init__(model=model, vector_dimension=vector_dimension) self.base_url = base_url or "http://localhost:11434"