戒酒的李白
43 lines
1.5 KiB
Python
43 lines
1.5 KiB
Python
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_
|