test_vectorizers.py - class2 - tarasovxx/AAA

test_vectorizers.py
116 строк4 KB
Tarasov Artyom
hw4: done
16 янв 2025, 21:11
16 янв 2025, 21:11bfe6d8c
from count_vectorizer import CountVectorizer
from tfidf_transformer import TfidfTransformer
from tfidf_vectorizer import TfidfVectorizer


def test_count_vectorizer():
    corpus = [
        'Crock Pot Pasta Never boil pasta again',
        'Pasta Pomodoro Fresh ingredients Parmesan to taste'
    ]
    vectorizer = CountVectorizer()
    count_matrix = vectorizer.fit_transform(corpus)
    feature_names = vectorizer.get_feature_names()

    expected_features = [
        'again', 'boil', 'crock', 'fresh', 'ingredients',
        'never', 'parmesan', 'pasta', 'pomodoro', 'pot', 'taste', 'to'
    ]
    expected_count_matrix = [
        [1, 1, 1, 0, 0, 1, 0, 2, 0, 1, 0, 0],
        [0, 0, 0, 1, 1, 0, 1, 1, 1, 0, 1, 1]
    ]

    assert feature_names == expected_features, (
        f"Ожидался словарь {expected_features}, но получен {feature_names}"
    )
    assert count_matrix == expected_count_matrix, (
        f"Ожидалась матрица {expected_count_matrix}, но получена {count_matrix}"
    )

    print("Тест test_count_vectorizer пройден!")


def test_tfidf_transformer():
    """
    Проверяем TfidfTransformer:
      1) Считаем tf
      2) Считаем idf
      3) Проверяем tf-idf
    На примере небольшой count_matrix.
    """
    count_matrix = [
        [1, 1, 2, 1, 1, 1, 0, 0, 0, 0, 0, 0],
        [0, 0, 1, 0, 0, 0, 1, 1, 1, 1, 1, 1]
    ]
    transformer = TfidfTransformer()
    tfidf_matrix = transformer.fit_transform(count_matrix)

    # Проверяем (пример из задания) idf:
    # Считаем, что если слово встретилось в 1 документе из 2, doc_freq=1,
    # то idf ~ ln((2+1)/(1+1)) + 1 = ln(3/2)+1 ~ ln(1.5)+1 ~ 0.405360 + 1 = 1.405360
    # Если слово встретилось в 2 документах, doc_freq=2,
    # idf ~ ln((2+1)/(2+1)) + 1 = ln(1) + 1 = 1.
    # Приблизительно ~ 1.4 и 1.0.
    # Проверим, что это сходится для первого и третьего слова (например).

    idf_values = transformer.get_idf()
    assert round(idf_values[0], 1) == 1.4, (
        f"Ожидался idf ~ 1.4, получен {idf_values[0]}"
    )
    assert round(idf_values[2], 1) == 1.0, (
        f"Ожидался idf ~ 1.0, получен {idf_values[2]}"
    )

    # Проверим несколько значений tf-idf на первом документе:
    # 'pasta' -> freq=2, sum=1+1+2+1+1+1=7 => tf=2/7=0.2857..., idf=1 => tf-idf ~ 0.286
    # 'crock' -> freq=1, tf=1/7=0.1428..., idf~1.405 => tf-idf ~ 0.2 (округляя до 0.1)
    # Для упрощения не будем делать точный assert для всего массива
    # (т.к. могут быть небольшие расхождения в float),
    # но проверим, что значения разумно близки.

    doc1_tfidf = tfidf_matrix[0]
    pasta_tfidf = doc1_tfidf[2]
    assert abs(pasta_tfidf - 0.286) < 0.05, (
        f"Ожидался tf-idf для 'pasta' ~ 0.286, получен {pasta_tfidf}"
    )

    print("Тест test_tfidf_transformer пройден!")


def test_tfidf_vectorizer():
    """
    Проверяем TfidfVectorizer "с нуля":
      - Сначала будет CountVectorizer
      - Затем TfidfTransformer
    """
    corpus = [
        'Crock Pot Pasta Never boil pasta again',
        'Pasta Pomodoro Fresh ingredients Parmesan to taste'
    ]
    vectorizer = TfidfVectorizer()
    tfidf_matrix = vectorizer.fit_transform(corpus)
    feature_names = vectorizer.get_feature_names()

    assert len(tfidf_matrix) == 2, "Ожидается 2 строки (2 документа)."

    expected_features = [
        'again', 'boil', 'crock', 'fresh', 'ingredients',
        'never', 'parmesan', 'pasta', 'pomodoro', 'pot', 'taste', 'to'
    ]
    assert feature_names == expected_features, (
        f"Ожидался словарь {expected_features}, получен {feature_names}"
    )

    print("Тест test_tfidf_vectorizer пройден!")


def main():
    test_count_vectorizer()
    test_tfidf_transformer()
    test_tfidf_vectorizer()
    print("\nВсе тесты пройдены успешно!")


if __name__ == "__main__":
    main()