Fixing the bug when using Huggingface Models (#1877)

Co-authored-by: parshvadaftari <[email protected]>

mem0/embeddings/huggingface.py

@@ -28,4 +28,4 @@ def embed(self, text):
         Returns:
             list: The embedding vector.
         """
-        return self.model.encode(text)
+        return self.model.encode(text, convert_to_numpy = True).tolist()

tests/embeddings/test_huggingface_embeddings.py

@@ -1,5 +1,6 @@
 import pytest
 from unittest.mock import Mock, patch
+import numpy as np
 from mem0.embeddings.huggingface import HuggingFaceEmbedding
 from mem0.configs.embeddings.base import BaseEmbedderConfig
 
@@ -16,35 +17,32 @@ def test_embed_default_model(mock_sentence_transformer):
     config = BaseEmbedderConfig()
     embedder = HuggingFaceEmbedding(config)
 
-    mock_sentence_transformer.encode.return_value = [0.1, 0.2, 0.3]
+    mock_sentence_transformer.encode.return_value = np.array([0.1, 0.2, 0.3])
     result = embedder.embed("Hello world")
 
-    mock_sentence_transformer.encode.assert_called_once_with("Hello world")
-
+    mock_sentence_transformer.encode.assert_called_once_with("Hello world", convert_to_numpy=True)
     assert result == [0.1, 0.2, 0.3]
 
 
 def test_embed_custom_model(mock_sentence_transformer):
     config = BaseEmbedderConfig(model="paraphrase-MiniLM-L6-v2")
     embedder = HuggingFaceEmbedding(config)
 
-    mock_sentence_transformer.encode.return_value = [0.4, 0.5, 0.6]
+    mock_sentence_transformer.encode.return_value = np.array([0.4, 0.5, 0.6])
     result = embedder.embed("Custom model test")
 
-    mock_sentence_transformer.encode.assert_called_once_with("Custom model test")
-
+    mock_sentence_transformer.encode.assert_called_once_with("Custom model test", convert_to_numpy=True)
     assert result == [0.4, 0.5, 0.6]
 
 
 def test_embed_with_model_kwargs(mock_sentence_transformer):
     config = BaseEmbedderConfig(model="all-MiniLM-L6-v2", model_kwargs={"device": "cuda"})
     embedder = HuggingFaceEmbedding(config)
 
-    mock_sentence_transformer.encode.return_value = [0.7, 0.8, 0.9]
+    mock_sentence_transformer.encode.return_value = np.array([0.7, 0.8, 0.9])
     result = embedder.embed("Test with device")
 
-    mock_sentence_transformer.encode.assert_called_once_with("Test with device")
-
+    mock_sentence_transformer.encode.assert_called_once_with("Test with device", convert_to_numpy=True)
     assert result == [0.7, 0.8, 0.9]
 
 
@@ -62,10 +60,10 @@ def test_embed_with_custom_embedding_dims(mock_sentence_transformer):
     config = BaseEmbedderConfig(model="all-mpnet-base-v2", embedding_dims=768)
     embedder = HuggingFaceEmbedding(config)
 
-    mock_sentence_transformer.encode.return_value = [1.0, 1.1, 1.2]
+    mock_sentence_transformer.encode.return_value = np.array([1.0, 1.1, 1.2])
     result = embedder.embed("Custom embedding dims")
 
-    mock_sentence_transformer.encode.assert_called_once_with("Custom embedding dims")
+    mock_sentence_transformer.encode.assert_called_once_with("Custom embedding dims", convert_to_numpy=True)
 
     assert embedder.config.embedding_dims == 768