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