langchain/tests/integration_tests/embeddings/test_mosaicml.py
Daniel King de6e6c764e
Add MosaicML inference endpoints (#4607)
# Add MosaicML inference endpoints
This PR adds support in langchain for MosaicML inference endpoints. We
both serve a select few open source models, and allow customers to
deploy their own models using our inference service. Docs are here
(https://docs.mosaicml.com/en/latest/inference.html), and sign up form
is here (https://forms.mosaicml.com/demo?utm_source=langchain). I'm not
intimately familiar with the details of langchain, or the contribution
process, so please let me know if there is anything that needs fixing or
this is the wrong way to submit a new integration, thanks!

I'm also not sure what the procedure is for integration tests. I have
tested locally with my api key.

## Who can review?
@hwchase17

---------

Co-authored-by: Harrison Chase <hw.chase.17@gmail.com>
2023-05-23 15:59:08 -07:00

59 lines
2.0 KiB
Python

"""Test mosaicml embeddings."""
from langchain.embeddings.mosaicml import MosaicMLInstructorEmbeddings
def test_mosaicml_embedding_documents() -> None:
"""Test MosaicML embeddings."""
documents = ["foo bar"]
embedding = MosaicMLInstructorEmbeddings()
output = embedding.embed_documents(documents)
assert len(output) == 1
assert len(output[0]) == 768
def test_mosaicml_embedding_documents_multiple() -> None:
"""Test MosaicML embeddings with multiple documents."""
documents = ["foo bar", "bar foo", "foo"]
embedding = MosaicMLInstructorEmbeddings()
output = embedding.embed_documents(documents)
assert len(output) == 3
assert len(output[0]) == 768
assert len(output[1]) == 768
assert len(output[2]) == 768
def test_mosaicml_embedding_query() -> None:
"""Test MosaicML embeddings of queries."""
document = "foo bar"
embedding = MosaicMLInstructorEmbeddings()
output = embedding.embed_query(document)
assert len(output) == 768
def test_mosaicml_embedding_endpoint() -> None:
"""Test MosaicML embeddings with a different endpoint"""
documents = ["foo bar"]
embedding = MosaicMLInstructorEmbeddings(
endpoint_url="https://models.hosted-on.mosaicml.hosting/instructor-xl/v1/predict"
)
output = embedding.embed_documents(documents)
assert len(output) == 1
assert len(output[0]) == 768
def test_mosaicml_embedding_query_instruction() -> None:
"""Test MosaicML embeddings with a different query instruction."""
document = "foo bar"
embedding = MosaicMLInstructorEmbeddings(query_instruction="Embed this query:")
output = embedding.embed_query(document)
assert len(output) == 768
def test_mosaicml_embedding_document_instruction() -> None:
"""Test MosaicML embeddings with a different query instruction."""
documents = ["foo bar"]
embedding = MosaicMLInstructorEmbeddings(embed_instruction="Embed this document:")
output = embedding.embed_documents(documents)
assert len(output) == 1
assert len(output[0]) == 768