Refactored embedding models and Updated documentation for Gemini model

mem0ai · Sep 30, 2024 · 61c703f · 61c703f
1 parent f324462
commit 61c703f
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 20 deletions.
diff --git a/docs/components/embedders/models/gemini.mdx b/docs/components/embedders/models/gemini.mdx
@@ -16,16 +16,9 @@ config = {
     "embedder": {
         "provider": "gemini",
         "config": {
-            "model": "models/text-embedding-004"
+            "model": "models/text-embedding-004",
         }
-    },
-    "vector_store": {
-        "provider": "qdrant",
-        "config": {
-            "collection_name": "test",
-            "embedding_model_dims": 768,
-        }
-    },
+    }
 }
 
 m = Memory.from_config(config)
@@ -39,3 +32,5 @@ Here are the parameters available for configuring Gemini embedder:
 | Parameter | Description | Default Value |
 | --- | --- | --- |
 | `model` | The name of the embedding model to use | `models/text-embedding-004` |
+| `embedding_dims` | Dimensions of the embedding model | `768` |
+| `api_key` | The Gemini API key | `None` |
diff --git a/mem0/embeddings/gemini.py b/mem0/embeddings/gemini.py
@@ -9,10 +9,13 @@
 class GoogleGenAIEmbedding(EmbeddingBase):
     def __init__(self, config: Optional[BaseEmbedderConfig] = None):
         super().__init__(config)
-        if self.config.model is None:
-            self.config.model = "models/text-embedding-004" # embedding-dim = 768
+
+        self.config.model = self.config.model or "models/text-embedding-004"
+        self.config.embedding_dims = self.config.embedding_dims or 768
 
-        genai.configure(api_key=self.config.api_key or os.getenv("GOOGLE_API_KEY"))
+        api_key = self.config.api_key or os.getenv("GOOGLE_API_KEY")
+
+        genai.configure(api_key=api_key)
 
     def embed(self, text):
         """

diff --git a/mem0/embeddings/huggingface.py b/mem0/embeddings/huggingface.py
@@ -10,13 +10,11 @@ class HuggingFaceEmbedding(EmbeddingBase):
     def __init__(self, config: Optional[BaseEmbedderConfig] = None):
         super().__init__(config)
 
-        if self.config.model is None:
-            self.config.model = "multi-qa-MiniLM-L6-cos-v1"
+        self.config.model = self.config.model or "multi-qa-MiniLM-L6-cos-v1"
 
         self.model = SentenceTransformer(self.config.model, **self.config.model_kwargs)
 
-        if self.config.embedding_dims is None:
-            self.config.embedding_dims = self.model.get_sentence_embedding_dimension()
+        self.config.embedding_dims = self.config.embedding_dims  or self.model.get_sentence_embedding_dimension()
 
     def embed(self, text):
         """

diff --git a/mem0/embeddings/ollama.py b/mem0/embeddings/ollama.py
@@ -25,10 +25,8 @@ class OllamaEmbedding(EmbeddingBase):
     def __init__(self, config: Optional[BaseEmbedderConfig] = None):
         super().__init__(config)
 
-        if not self.config.model:
-            self.config.model = "nomic-embed-text"
-        if not self.config.embedding_dims:
-            self.config.embedding_dims = 512
+        self.config.model = self.config.model or "nomic-embed-text"
+        self.config.embedding_dims = self.config.embedding_dims or 512
 
         self.client = Client(host=self.config.ollama_base_url)
         self._ensure_model_exists()