Add test for svm retriever (#7768)

# What
- This is to add unit test for svm retriever.

---------

Co-authored-by: Bagatur <baskaryan@gmail.com>
This commit is contained in:
shibuiwilliam 2023-07-19 01:57:24 +09:00 committed by GitHub
parent 69b9db2b5e
commit 177baef3a1
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -0,0 +1,27 @@
import pytest
from langchain.embeddings import FakeEmbeddings
from langchain.retrievers.svm import SVMRetriever
from langchain.schema import Document
class TestSVMRetriever:
@pytest.mark.requires("sklearn")
def test_from_texts(self) -> None:
input_texts = ["I have a pen.", "Do you have a pen?", "I have a bag."]
svm_retriever = SVMRetriever.from_texts(
texts=input_texts, embeddings=FakeEmbeddings(size=100)
)
assert len(svm_retriever.texts) == 3
@pytest.mark.requires("sklearn")
def test_from_documents(self) -> None:
input_docs = [
Document(page_content="I have a pen.", metadata={"foo": "bar"}),
Document(page_content="Do you have a pen?"),
Document(page_content="I have a bag."),
]
svm_retriever = SVMRetriever.from_documents(
documents=input_docs, embeddings=FakeEmbeddings(size=100)
)
assert len(svm_retriever.texts) == 3