test_embeddings.py
2 KB
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
57
58
59
60
61
62
63
64
65
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"])