| | 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_ |
| |
|
