klas/t6_mem0
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[ Refactored ] embedding models and [ Update ] documentation for Gemini model (#1931)

Browse Source
This commit is contained in:
Parshva Daftari
2024-10-03 21:30:46 +05:30
committed by GitHub
parent 395af18d88
commit c09c4926a7
4 changed files with 14 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
docs/components/embedders/models/gemini.mdx
View File

@@ -16,16 +16,9 @@ config = {
"embedder": { "embedder": {
"provider": "gemini", "provider": "gemini",
"config": { "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) m = Memory.from_config(config)
@@ -39,3 +32,5 @@ Here are the parameters available for configuring Gemini embedder:
| Parameter | Description | Default Value | | Parameter | Description | Default Value |
| --- | --- | --- | | --- | --- | --- |
| `model` | The name of the embedding model to use | `models/text-embedding-004` | | `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` |

9
mem0/embeddings/gemini.py
View File

@@ -9,10 +9,13 @@ from mem0.embeddings.base import EmbeddingBase
class GoogleGenAIEmbedding(EmbeddingBase): class GoogleGenAIEmbedding(EmbeddingBase):
def __init__(self, config: Optional[BaseEmbedderConfig] = None): def __init__(self, config: Optional[BaseEmbedderConfig] = None):
super().__init__(config) 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): def embed(self, text):
""" """

6
mem0/embeddings/huggingface.py
View File

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

6
mem0/embeddings/ollama.py
View File

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