Add support for configurable embedding model (#1627)

Co-authored-by: Dev Khant <devkhant24@gmail.com>

mem0/embeddings/azure_openai.py

@@ -0,0 +1,36 @@
+from typing import Optional
+
+from openai import AzureOpenAI
+
+from mem0.configs.embeddings.base import BaseEmbedderConfig
+from mem0.embeddings.base import EmbeddingBase
+
+class AzureOpenAIEmbedding(EmbeddingBase):
+    def __init__(self, config: Optional[BaseEmbedderConfig] = None):
+        super().__init__(config)
+        
+        if self.config.model is None:
+            self.config.model = "text-embedding-3-small"
+        if self.config.embedding_dims is None:
+            self.config.embedding_dims = 1536
+        self.client = AzureOpenAI()
+
+    def embed(self, text):
+        """
+        Get the embedding for the given text using OpenAI.
+
+        Args:
+            text (str): The text to embed.
+
+        Returns:
+            list: The embedding vector.
+        """
+        text = text.replace("\n", " ")
+        return (
+            self.client.embeddings.create(
+                input=[text], 
+                model=self.config.model
+            )
+            .data[0]
+            .embedding
+        )

mem0/embeddings/base.py

@@ -3,19 +3,18 @@ from abc import ABC, abstractmethod
 
 from mem0.configs.embeddings.base import BaseEmbedderConfig
 
-
 class EmbeddingBase(ABC):
-    def __init__(self, config: Optional[BaseEmbedderConfig] = None):
-        """Initialize a base LLM class
+    """Initialized a base embedding class
 
-        :param config: Embedder configuration option class, defaults to None
-        :type config: Optional[BaseEmbedderConfig], optional
-        """
+    :param config: Embedding configuration option class, defaults to None
+    :type config: Optional[BaseEmbedderConfig], optional
+    """
+    def __init__(self, config: Optional[BaseEmbedderConfig] = None):
         if config is None:
             self.config = BaseEmbedderConfig()
         else:
             self.config = config
-
+            
     @abstractmethod
     def embed(self, text):
         """

mem0/embeddings/configs.py

@@ -9,13 +9,14 @@ class EmbedderConfig(BaseModel):
         default="openai",
     )
     config: Optional[dict] = Field(
-        description="Configuration for the specific embedding model", default=None
+        description="Configuration for the specific embedding model", 
+        default={}
     )
 
     @field_validator("config")
     def validate_config(cls, v, values):
         provider = values.data.get("provider")
-        if provider in ["openai", "ollama"]:
+        if provider in ["openai", "ollama", "huggingface", "azure_openai"]:
             return v
         else:
             raise ValueError(f"Unsupported embedding provider: {provider}")

mem0/embeddings/huggingface.py

@@ -1,11 +1,27 @@
-from mem0.embeddings.base import EmbeddingBase
+from typing import Optional
+
 from sentence_transformers import SentenceTransformer
 
+from mem0.configs.embeddings.base import BaseEmbedderConfig
+from mem0.embeddings.base import EmbeddingBase
+
 
 class HuggingFaceEmbedding(EmbeddingBase):
-    def __init__(self, model_name="multi-qa-MiniLM-L6-cos-v1"):
-        self.model = SentenceTransformer(model_name)
+    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.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()
+
+    
     def embed(self, text):
         """
         Get the embedding for the given text using Hugging Face.

mem0/embeddings/openai.py

@@ -1,10 +1,10 @@
 from typing import Optional
+
 from openai import OpenAI
 
 from mem0.configs.embeddings.base import BaseEmbedderConfig
 from mem0.embeddings.base import EmbeddingBase
 
-
 class OpenAIEmbedding(EmbeddingBase):
     def __init__(self, config: Optional[BaseEmbedderConfig] = None):
         super().__init__(config)
@@ -28,7 +28,10 @@ class OpenAIEmbedding(EmbeddingBase):
         """
         text = text.replace("\n", " ")
         return (
-            self.client.embeddings.create(input=[text], model=self.config.model)
+            self.client.embeddings.create(
+                input=[text], 
+                model=self.config.model
+            )
             .data[0]
             .embedding
         )