diff --git a/bertopic/_bertopic.py b/bertopic/_bertopic.py index cfafb58a..bb83b2a7 100644 --- a/bertopic/_bertopic.py +++ b/bertopic/_bertopic.py @@ -2188,6 +2188,7 @@ def delete_topics( """ check_is_fitted(self) + original_unique_topics = sorted(set(self.topics_)) topics_df = pd.DataFrame({"Topic": self.topics_}) # Check if -1 exists in the current topics @@ -2215,10 +2216,6 @@ 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 - if hasattr(self, "custom_labels_") and self.custom_labels_ is not None: - self.custom_labels_[-1] = "" - # Initialize ctfidf model diagonal for -1 topic (ones) if it exists if hasattr(self, "ctfidf_model") and self.ctfidf_model is not None: n_features = self.ctfidf_model._idf_diag.shape[1] @@ -2268,14 +2265,23 @@ def delete_topics( } self.topic_aspects_ = new_aspects - # Update custom labels if they exist + # Update custom labels if they exist. `.custom_labels_` is a list aligned with + # sorted unique topic IDs (see `set_topic_labels`), not a dict. 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 isinstance(self.custom_labels_, dict): + labels_by_topic = dict(self.custom_labels_) + else: + labels_by_topic = dict(zip(original_unique_topics, self.custom_labels_)) + + remapped_labels = {} + for old_topic, label in labels_by_topic.items(): + if old_topic in topics_to_delete: + continue + new_topic = -1 if old_topic == -1 else final_mapping[old_topic] + remapped_labels[new_topic] = label + + new_unique_topics = sorted(set(self.topics_)) + self.custom_labels_ = [remapped_labels.get(topic, "") for topic in new_unique_topics] # Update topic representations new_representations = { diff --git a/tests/test_reduction/test_delete.py b/tests/test_reduction/test_delete.py index 188e1ffb..7b2805c6 100644 --- a/tests/test_reduction/test_delete.py +++ b/tests/test_reduction/test_delete.py @@ -57,3 +57,26 @@ 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"), + ], +) +def test_delete_keeps_custom_labels(model, request): + topic_model = copy.deepcopy(request.getfixturevalue(model)) + labels = topic_model.generate_topic_labels() + topic_model.set_topic_labels(labels) + + remaining = [topic for topic in sorted(set(topic_model.topics_)) if topic != -1] + topic_to_delete = remaining[0] + + topic_model.delete_topics([topic_to_delete]) + + unique_topics = sorted(set(topic_model.topics_)) + assert isinstance(topic_model.custom_labels_, list) + assert len(topic_model.custom_labels_) == len(unique_topics)