Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 1 commit into from
Oct 3, 2024
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 4 additions & 9 deletions docs/components/embedders/models/gemini.mdx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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` |
9 changes: 6 additions & 3 deletions mem0/embeddings/gemini.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
"""
Expand Down
6 changes: 2 additions & 4 deletions mem0/embeddings/huggingface.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
"""
Expand Down
6 changes: 2 additions & 4 deletions mem0/embeddings/ollama.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
Loading