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

Jump to bottom

Fixing the bug when using Huggingface Models #1877

Merged
merged 5 commits into from
Sep 25, 2024
Merged

Fixing the bug when using Huggingface Models #1877

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
6b9ddf0
Fixing the bug when using Huggingface Models
Sep 17, 2024
606ae13
Merge branch 'mem0ai:main' into parshva/huggingface-fix
parshvadaftari Sep 21, 2024
c577010
Linting the code for better readability
parshvadaftari Sep 25, 2024
f810824
Merge branch 'parshva/huggingface-fix' of github.com:parshvadaftari/m…
parshvadaftari Sep 25, 2024
bf8467d
Test script for huggingface embeddings
parshvadaftari Sep 25, 2024
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
2 changes: 1 addition & 1 deletion mem0/embeddings/huggingface.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
20 changes: 9 additions & 11 deletions tests/embeddings/test_huggingface_embeddings.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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

Expand All @@ -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]


Expand All @@ -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

Expand Down
Loading