forked from Archives/langchain
de6e6c764e
# 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>
59 lines
2.0 KiB
Python
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
|