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

Add embedder docs and config changes (#1684)

Browse Source
This commit is contained in:
Dev Khant
2024-08-12 16:09:01 +05:30
committed by GitHub
parent 464a188662
commit b245309242
12 changed files with 130 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
mem0/embeddings/ollama.py
View File

@@ -18,7 +18,7 @@ class OllamaEmbedding(EmbeddingBase):
if not self.config.embedding_dims:
self.config.embedding_dims=512
self.client = Client(host=self.config.base_url)
self.client = Client(host=self.config.ollama_base_url)
self._ensure_model_exists()
def _ensure_model_exists(self):

10
mem0/embeddings/openai.py
View File

@@ -1,3 +1,4 @@
import os
from typing import Optional
from openai import OpenAI
@@ -9,12 +10,11 @@ class OpenAIEmbedding(EmbeddingBase):
def __init__(self, config: Optional[BaseEmbedderConfig] = None):
super().__init__(config)
if not self.config.model:
self.config.model="text-embedding-3-small"
if not self.config.embedding_dims:
self.config.embedding_dims=1536
self.config.model = self.config.model or "text-embedding-3-small"
self.config.embedding_dims = self.config.embedding_dims or 1536
self.client = OpenAI()
api_key = os.getenv("OPENAI_API_KEY") or self.config.api_key
self.client = OpenAI(api_key=api_key)
def embed(self, text):
"""