langchain/libs/community/tests/unit_tests/retrievers/test_tfidf.py

import os
from datetime import datetime
from tempfile import TemporaryDirectory

import pytest
from langchain_core.documents import Document

from langchain_community.retrievers.tfidf import TFIDFRetriever


@pytest.mark.requires("sklearn")
def test_from_texts() -> None:
    input_texts = ["I have a pen.", "Do you have a pen?", "I have a bag."]
    tfidf_retriever = TFIDFRetriever.from_texts(texts=input_texts)
    assert len(tfidf_retriever.docs) == 3
    assert tfidf_retriever.tfidf_array.toarray().shape == (3, 5)


@pytest.mark.requires("sklearn")
def test_from_texts_with_tfidf_params() -> None:
    input_texts = ["I have a pen.", "Do you have a pen?", "I have a bag."]
    tfidf_retriever = TFIDFRetriever.from_texts(
        texts=input_texts, tfidf_params={"min_df": 2}
    )
    # should count only multiple words (have, pan)
    assert tfidf_retriever.tfidf_array.toarray().shape == (3, 2)


@pytest.mark.requires("sklearn")
def test_from_documents() -> None:
    input_docs = [
        Document(page_content="I have a pen."),
        Document(page_content="Do you have a pen?"),
        Document(page_content="I have a bag."),
    ]
    tfidf_retriever = TFIDFRetriever.from_documents(documents=input_docs)
    assert len(tfidf_retriever.docs) == 3
    assert tfidf_retriever.tfidf_array.toarray().shape == (3, 5)


@pytest.mark.requires("sklearn")
def test_save_local_load_local() -> None:
    input_texts = ["I have a pen.", "Do you have a pen?", "I have a bag."]
    tfidf_retriever = TFIDFRetriever.from_texts(texts=input_texts)

    file_name = "tfidf_vectorizer"
    temp_timestamp = datetime.utcnow().strftime("%Y%m%d-%H%M%S")
    with TemporaryDirectory(suffix="_" + temp_timestamp + "/") as temp_folder:
        tfidf_retriever.save_local(
            folder_path=temp_folder,
            file_name=file_name,
        )
        assert os.path.exists(os.path.join(temp_folder, f"{file_name}.joblib"))
        assert os.path.exists(os.path.join(temp_folder, f"{file_name}.pkl"))

        loaded_tfidf_retriever = TFIDFRetriever.load_local(
            folder_path=temp_folder,
            file_name=file_name,
        )
    assert len(loaded_tfidf_retriever.docs) == 3
    assert loaded_tfidf_retriever.tfidf_array.toarray().shape == (3, 5)