Jelajahi Sumber

Add vector dimension to Ollama embedder (#1435)

Dev Khant 1 tahun lalu
induk
melakukan
b43a116b3c

+ 2 - 1
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"

+ 2 - 1
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"