Add support for configurable embedding model (#1627)
Co-authored-by: Dev Khant <devkhant24@gmail.com>
This commit is contained in:
@@ -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.
|
||||
|
||||
Reference in New Issue
Block a user