diff --git a/langchain/embeddings/huggingface.py b/langchain/embeddings/huggingface.py new file mode 100644 index 0000000000..cf59a300d4 --- /dev/null +++ b/langchain/embeddings/huggingface.py @@ -0,0 +1,66 @@ +"""Wrapper around HuggingFace embedding models.""" +from typing import Any, List + +from pydantic import BaseModel, Extra + +from langchain.embeddings.base import Embeddings + + +class HuggingFaceEmbeddings(BaseModel, Embeddings): + """Wrapper around sentence_transformers embedding models. + + To use, you should have the ``sentence_transformers`` python package installed. + + Example: + .. code-block:: python + + from langchain.embeddings import HuggingFaceEmbeddings + huggingface = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2") + """ + + client: Any #: :meta private: + model_name: str = "sentence-transformers/all-mpnet-base-v2" + """Model name to use.""" + + def __init__(self, **kwargs: Any): + super().__init__(**kwargs) + try: + import sentence_transformers + + self.client = sentence_transformers.SentenceTransformer(self.model_name) + except ImportError: + raise ValueError( + "Could not import sentence_transformers python package. " + "Please install it with `pip install sentence_transformers`." + ) + + class Config: + """Configuration for this pydantic object.""" + + extra = Extra.forbid + + def embed_documents(self, texts: List[str]) -> List[List[float]]: + """Computes doc embeddings using a HuggingFace transformer model + + Args: + texts: The list of texts to embed. + + Returns: + List of embeddings, one for each text. + """ + texts = list(map(lambda x: x.replace("\n", " "), texts)) + embeddings = self.client.encode(texts) + return embeddings + + def embed_query(self, text: str) -> List[float]: + """Computes query embeddings using a HuggingFace transformer model + + Args: + text: The text to embed. + + Returns: + Embeddings for the text. + """ + text = text.replace("\n", " ") + embedding = self.client.encode(text) + return embedding diff --git a/requirements.txt b/requirements.txt index 1dd1e8a215..c60ee0cd91 100644 --- a/requirements.txt +++ b/requirements.txt @@ -13,5 +13,6 @@ playwright wikipedia huggingface_hub faiss +sentence_transformers # For development jupyter diff --git a/tests/integration_tests/embeddings/test_huggingface.py b/tests/integration_tests/embeddings/test_huggingface.py new file mode 100644 index 0000000000..5c031ee045 --- /dev/null +++ b/tests/integration_tests/embeddings/test_huggingface.py @@ -0,0 +1,19 @@ +"""Test huggingface embeddings.""" +from langchain.embeddings.huggingface import HuggingFaceEmbeddings + + +def test_huggingface_embedding_documents() -> None: + """Test huggingface embeddings.""" + documents = ["foo bar"] + embedding = HuggingFaceEmbeddings() + output = embedding.embed_documents(documents) + assert len(output) == 1 + assert len(output[0]) == 768 + + +def test_huggingface_embedding_query() -> None: + """Test huggingface embeddings.""" + document = "foo bar" + embedding = HuggingFaceEmbeddings() + output = embedding.embed_query(document) + assert len(output) == 768