BERTopic/tests/test_sub_models/test_embeddings.py

import copy
import pytest
import numpy as np
from bertopic import BERTopic
from sklearn.metrics.pairwise import cosine_similarity


@pytest.mark.parametrize('model', [('kmeans_pca_topic_model'),
                                   ('base_topic_model'),
                                   ('custom_topic_model'),
                                   ('merged_topic_model'),
                                   ('reduced_topic_model'),
                                   ('online_topic_model')])
def test_extract_embeddings(model, request):
    topic_model = copy.deepcopy(request.getfixturevalue(model))
    single_embedding = topic_model._extract_embeddings("a document")
    multiple_embeddings = topic_model._extract_embeddings(["something different", "another document"])
    sim_matrix = cosine_similarity(single_embedding, multiple_embeddings)[0]

    assert single_embedding.shape[0] == 1
    assert single_embedding.shape[1] == 384
    assert np.min(single_embedding) > -5
    assert np.max(single_embedding) < 5

    assert multiple_embeddings.shape[0] == 2
    assert multiple_embeddings.shape[1] == 384
    assert np.min(multiple_embeddings) > -5
    assert np.max(multiple_embeddings) < 5

    assert sim_matrix[0] < 0.5
    assert sim_matrix[1] > 0.5


@pytest.mark.parametrize('model', [('kmeans_pca_topic_model'),
                                   ('base_topic_model'),
                                   ('custom_topic_model'),
                                   ('merged_topic_model'),
                                   ('reduced_topic_model'),
                                   ('online_topic_model')])
def test_extract_embeddings_compare(model, embedding_model, request):
    docs = ["some document"]
    topic_model = copy.deepcopy(request.getfixturevalue(model))
    bertopic_embeddings = topic_model._extract_embeddings(docs)

    assert isinstance(bertopic_embeddings, np.ndarray)
    assert bertopic_embeddings.shape == (1, 384)

    sentence_embeddings = embedding_model.encode(docs, show_progress_bar=False)
    assert np.array_equal(bertopic_embeddings, sentence_embeddings)


def test_extract_incorrect_embeddings():
    with pytest.raises(ValueError):
        model = BERTopic(language="Unknown language")
        model.fit(["some document"])