test_merge.py
1.5 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
import copy
import pytest
@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_merge(model, documents, request):
topic_model = copy.deepcopy(request.getfixturevalue(model))
nr_topics = len(set(topic_model.topics_))
topics_to_merge = [1, 2]
topic_model.merge_topics(documents, topics_to_merge)
mappings = topic_model.topic_mapper_.get_mappings(list(topic_model.hdbscan_model.labels_))
mapped_labels = [mappings[label] for label in topic_model.hdbscan_model.labels_]
assert nr_topics == len(set(topic_model.topics_)) + 1
assert topic_model.get_topic_info().Count.sum() == len(documents)
if model == "online_topic_model":
assert mapped_labels == topic_model.topics_[950:]
else:
assert mapped_labels == topic_model.topics_
topics_to_merge = [1, 2]
topic_model.merge_topics(documents, topics_to_merge)
mappings = topic_model.topic_mapper_.get_mappings(list(topic_model.hdbscan_model.labels_))
mapped_labels = [mappings[label] for label in topic_model.hdbscan_model.labels_]
assert nr_topics == len(set(topic_model.topics_)) + 2
assert topic_model.get_topic_info().Count.sum() == len(documents)
if model == "online_topic_model":
assert mapped_labels == topic_model.topics_[950:]
else:
assert mapped_labels == topic_model.topics_