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