diff --git a/bertopic/_bertopic.py b/bertopic/_bertopic.py index cfafb58a..480944df 100644 --- a/bertopic/_bertopic.py +++ b/bertopic/_bertopic.py @@ -2193,6 +2193,10 @@ def delete_topics( # Check if -1 exists in the current topics had_outliers = -1 in set(self.topics_) + # `custom_labels_` is a list ordered by topic rather than a mapping, so keep + # track of the topics it currently refers to in order to remap it below + custom_labels_topics = sorted(set(self.topics_)) + # If adding -1 for the first time, initialize its attributes if not had_outliers and any(topic in topics_to_delete for topic in self.topics_): # Initialize c-TF-IDF for -1 topic (zeros) @@ -2215,9 +2219,11 @@ def delete_topics( outlier_image = np.zeros((1, self.representative_images_.shape[1])) self.representative_images_ = np.vstack([outlier_image, self.representative_images_]) - # Initialize custom labels for -1 topic if they exist + # Initialize custom labels for -1 topic if they exist. Topic -1 sorts + # first, so its label is prepended rather than assigned to index -1. if hasattr(self, "custom_labels_") and self.custom_labels_ is not None: - self.custom_labels_[-1] = "" + self.custom_labels_.insert(0, "") + custom_labels_topics.insert(0, -1) # Initialize ctfidf model diagonal for -1 topic (ones) if it exists if hasattr(self, "ctfidf_model") and self.ctfidf_model is not None: @@ -2270,12 +2276,13 @@ def delete_topics( # Update custom labels if they exist if hasattr(self, "custom_labels_") and self.custom_labels_ is not None: - new_labels = { - (final_mapping[old_topic] if old_topic != -1 else -1): label - for old_topic, label in self.custom_labels_.items() - if old_topic not in topics_to_delete - } - self.custom_labels_ = new_labels + if len(self.custom_labels_) == len(custom_labels_topics): + new_labels = { + (final_mapping[old_topic] if old_topic != -1 else -1): label + for old_topic, label in zip(custom_labels_topics, self.custom_labels_) + if old_topic not in topics_to_delete + } + self.custom_labels_ = [new_labels[topic] for topic in sorted(new_labels)] # Update topic representations new_representations = { diff --git a/tests/test_reduction/test_delete.py b/tests/test_reduction/test_delete.py index 188e1ffb..594ccc1e 100644 --- a/tests/test_reduction/test_delete.py +++ b/tests/test_reduction/test_delete.py @@ -57,3 +57,43 @@ def test_delete(model, request): assert mapped_labels == topic_model.topics_[950:] else: assert mapped_labels == topic_model.topics_ + + +@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_with_custom_labels(model, request): + """Custom labels are a list, so deleting topics must remap them positionally.""" + topic_model = copy.deepcopy(request.getfixturevalue(model)) + + # Label every topic after itself so misalignment is detectable + original_topics = sorted(set(topic_model.topics_)) + topic_model.set_topic_labels([f"label of topic {topic}" for topic in original_topics]) + labels_before = dict(zip(original_topics, topic_model.custom_labels_)) + + topics_to_delete = [topic for topic in original_topics if topic != -1][:2] + topic_model.delete_topics(topics_to_delete) + + remaining_topics = sorted(set(topic_model.topics_)) + assert isinstance(topic_model.custom_labels_, list) + assert len(topic_model.custom_labels_) == len(remaining_topics) + + # Every surviving topic keeps its own label, even though topics are renumbered + labels_after = dict(zip(remaining_topics, topic_model.custom_labels_)) + mappings = topic_model.topic_mapper_.get_mappings(original_topics=False) + for topic in original_topics: + if topic in topics_to_delete or topic == -1: + continue + assert labels_after[mappings[topic]] == labels_before[topic] + + # A newly created outlier topic gets an empty label rather than stealing one + if -1 not in original_topics: + assert labels_after[-1] == ""