File size: 2,270 Bytes
19b102a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 |
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"])
|