| 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"]) |
|
|
