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

9
mem0/embeddings/gemini.py
View File

@@ -9,10 +9,13 @@ from mem0.embeddings.base import EmbeddingBase
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):
"""

6
mem0/embeddings/huggingface.py
View File

@@ -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):
"""

6
mem0/embeddings/ollama.py
View File

@@ -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()