langchain/tests/integration_tests/embeddings/test_huggingface_hub.py
Bagatur b90e25f786
Add HuggingFace Hub Embeddings (#125)
Add support for calling HuggingFace embedding models
using the HuggingFaceHub Inference API. New class mirrors
the existing HuggingFaceHub LLM implementation. Currently
only supports 'sentence-transformers' models.

Closes #86
2022-11-27 00:24:59 -08:00

20 lines
615 B
Python

"""Test HuggingFaceHub embeddings."""
from langchain.embeddings import HuggingFaceHubEmbeddings
def test_huggingfacehub_embedding_documents() -> None:
"""Test huggingfacehub embeddings."""
documents = ["foo bar"]
embedding = HuggingFaceHubEmbeddings()
output = embedding.embed_documents(documents)
assert len(output) == 1
assert len(output[0]) == 768
def test_huggingfacehub_embedding_query() -> None:
"""Test huggingfacehub embeddings."""
document = "foo bar"
embedding = HuggingFaceHubEmbeddings()
output = embedding.embed_query(document)
assert len(output) == 768