戒酒的李白
60 lines
2.3 KiB
Python
60 lines
2.3 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_delete(model, request):
|
|
topic_model = copy.deepcopy(request.getfixturevalue(model))
|
|
nr_topics = len(set(topic_model.topics_))
|
|
length_documents = len(topic_model.topics_)
|
|
|
|
# First deletion
|
|
topics_to_delete = [1, 2]
|
|
topic_model.delete_topics(topics_to_delete)
|
|
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_]
|
|
|
|
if model == "online_topic_model" or model == "kmeans_pca_topic_model":
|
|
assert nr_topics == len(set(topic_model.topics_)) + 1
|
|
assert topic_model.get_topic_info().Count.sum() == length_documents
|
|
else:
|
|
assert nr_topics == len(set(topic_model.topics_)) + 2
|
|
assert topic_model.get_topic_info().Count.sum() == length_documents
|
|
|
|
if model == "online_topic_model":
|
|
assert mapped_labels == topic_model.topics_[950:]
|
|
else:
|
|
assert mapped_labels == topic_model.topics_
|
|
|
|
# Find two existing topics for second deletion
|
|
remaining_topics = sorted(list(set(topic_model.topics_)))
|
|
remaining_topics = [t for t in remaining_topics if t != -1] # Exclude outlier topic
|
|
topics_to_delete = remaining_topics[:2] # Take first two remaining topics
|
|
|
|
# Second deletion
|
|
topic_model.delete_topics(topics_to_delete)
|
|
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_]
|
|
|
|
if model == "online_topic_model" or model == "kmeans_pca_topic_model":
|
|
assert nr_topics == len(set(topic_model.topics_)) + 3
|
|
assert topic_model.get_topic_info().Count.sum() == length_documents
|
|
else:
|
|
assert nr_topics == len(set(topic_model.topics_)) + 4
|
|
assert topic_model.get_topic_info().Count.sum() == length_documents
|
|
|
|
if model == "online_topic_model":
|
|
assert mapped_labels == topic_model.topics_[950:]
|
|
else:
|
|
assert mapped_labels == topic_model.topics_
|