diff --git a/bertopic/_bertopic.py b/bertopic/_bertopic.py index cfafb58a..08c6702d 100644 --- a/bertopic/_bertopic.py +++ b/bertopic/_bertopic.py @@ -1038,6 +1038,7 @@ def hierarchical_topics( use_ctfidf: bool = True, linkage_function: Callable[[csr_matrix], np.ndarray] | None = None, distance_function: Callable[[csr_matrix], csr_matrix] | None = None, + use_representation_model: bool = False, ) -> pd.DataFrame: """Create a hierarchy of topics. @@ -1063,6 +1064,11 @@ def hierarchical_topics( non-negative values or condensed distance matrix of shape (n_samples * (n_samples - 1) / 2,) containing the upper triangular of the distance matrix. + use_representation_model: If True, after computing the hierarchy, run the + representation model (including aspects) on all parent + topics in a single batch to generate human-readable labels. + The ``Parent_Name`` column will contain the representation + model output instead of raw c-TF-IDF keyword concatenations. Returns: hierarchical_topics: A dataframe that contains a hierarchy of topics @@ -1126,6 +1132,10 @@ def hierarchical_topics( bow = self.vectorizer_model.transform(clean_documents) + # Collect per-parent data for batch relabeling + parent_c_tf_idf_rows = [] + parent_doc_selections = [] + # Extract clusters hier_topics = pd.DataFrame( columns=[ @@ -1160,6 +1170,11 @@ def hierarchical_topics( selection.Topic = 0 words_per_topic = self._extract_words_per_topic(words, selection, c_tf_idf, calculate_aspects=False) + # Save per-parent data for batch relabeling + if use_representation_model: + parent_c_tf_idf_rows.append(c_tf_idf) + parent_doc_selections.append(selection) + # Extract parent's name and ID parent_id = index + len(clusters) parent_name = "_".join([x[0] for x in words_per_topic[0]][:5]) @@ -1199,6 +1214,107 @@ def hierarchical_topics( ["Parent_ID", "Child_Left_ID", "Child_Right_ID"] ].astype(str) + # Batch relabeling with representation model + if use_representation_model and self.representation_model and parent_c_tf_idf_rows: + hier_topics = self._label_hierarchical_topics( + hier_topics, words, parent_c_tf_idf_rows, parent_doc_selections + ) + + return hier_topics + + def _label_hierarchical_topics( + self, + hier_topics: pd.DataFrame, + words: np.ndarray, + parent_c_tf_idf_rows: list, + parent_doc_selections: list, + ) -> pd.DataFrame: + """Relabel parent topics using the representation model in a single batch. + + Stacks all per-parent c-TF-IDF rows into one matrix, builds a combined + documents DataFrame with unique synthetic topic IDs, and calls + ``_extract_words_per_topic`` once so that both the Main model and aspect + models run on all parents together. + + Arguments: + hier_topics: The hierarchy DataFrame produced by the merge loop. + words: Feature names from the fitted vectorizer. + parent_c_tf_idf_rows: One c-TF-IDF row (sparse matrix) per parent. + parent_doc_selections: Document selection DataFrames per parent. + + Returns: + The updated ``hier_topics`` DataFrame with ``Parent_Name`` relabeled. + """ + from scipy.sparse import vstack as sparse_vstack + + n_parents = len(parent_c_tf_idf_rows) + batched_c_tf_idf = sparse_vstack(parent_c_tf_idf_rows, format="csr") + + # Build combined documents with unique topic IDs per parent + combined_docs_parts = [] + for parent_idx, selection in enumerate(parent_doc_selections): + part = selection.copy() + part["Topic"] = parent_idx + combined_docs_parts.append(part) + combined_docs = pd.concat(combined_docs_parts, ignore_index=True) + + # Save mutable state that _extract_words_per_topic overwrites + saved_topic_aspects = getattr(self, "topic_aspects_", {}) + saved_representative_docs = getattr(self, "representative_docs_", {}) + self.topic_aspects_ = {} + + try: + # Run representation model (Main + aspects) in one batched call + labeled_topics = self._extract_words_per_topic( + words, + combined_docs, + batched_c_tf_idf, + fine_tune_representation=True, + calculate_aspects=True, + ) + + parent_aspects = dict(self.topic_aspects_) + finally: + # Restore leaf-topic state even if _extract_words_per_topic raises + self.topic_aspects_ = saved_topic_aspects + self.representative_docs_ = saved_representative_docs + + # Determine the best label source (prefer non-Main aspect) + llm_aspect_name = None + if isinstance(self.representation_model, dict): + for aspect in self.representation_model: + if aspect != "Main": + llm_aspect_name = aspect + break + + # Build parent index to label mapping + parent_labels = {} + for parent_idx in range(n_parents): + label = None + if llm_aspect_name and llm_aspect_name in parent_aspects: + aspect_data = parent_aspects[llm_aspect_name].get(parent_idx, []) + if aspect_data: + label = str(aspect_data[0][0]) + if not label: + topic_words = labeled_topics.get(parent_idx, []) + if topic_words: + label = "_".join([w for w, _ in topic_words[:5]]) + if label: + parent_labels[parent_idx] = label + + # Update Parent_Name and propagate to child references + old_to_new = {} + for original_idx, new_name in parent_labels.items(): + if original_idx in hier_topics.index: + old_name = hier_topics.loc[original_idx, "Parent_Name"] + if old_name != new_name: + old_to_new[old_name] = new_name + hier_topics.loc[original_idx, "Parent_Name"] = new_name + + if old_to_new: + hier_topics["Child_Left_Name"] = hier_topics["Child_Left_Name"].replace(old_to_new) + hier_topics["Child_Right_Name"] = hier_topics["Child_Right_Name"].replace(old_to_new) + return hier_topics def approximate_distribution( diff --git a/docs/getting_started/hierarchicaltopics/hierarchicaltopics.md b/docs/getting_started/hierarchicaltopics/hierarchicaltopics.md index 8780b668..d42cb08d 100644 --- a/docs/getting_started/hierarchicaltopics/hierarchicaltopics.md +++ b/docs/getting_started/hierarchicaltopics/hierarchicaltopics.md @@ -35,6 +35,36 @@ hierarchical_topics = topic_model.hierarchical_topics(docs) The resulting `hierarchical_topics` is a dataframe in which merged topics are described. For example, if you would merge two topics, what would the topic representation of the new topic be? + +## **Rich labels for parent topics** + +By default, parent topics in the hierarchy are named using raw keyword +concatenation (e.g., `"safety_equipment_ppe_wear_worker"`). If you have a +representation model configured (e.g., KeyBERTInspired, LLM-based), you can +use it to generate rich labels for parent topics too: + +```python +from bertopic.representation import KeyBERTInspired + +representation_model = KeyBERTInspired() +topic_model = BERTopic(representation_model=representation_model) +topics, probs = topic_model.fit_transform(docs) + +# Rich labels for the hierarchy +hierarchical_topics = topic_model.hierarchical_topics( + docs, use_representation_model=True +) +``` + +This runs the representation model on all parent topics in a single batch +call after the hierarchy is built, so it adds only one round of inference +regardless of the number of merge steps. + +!!! note + When `representation_model` is `None`, setting `use_representation_model=True` + is a no-op — parent names default to keyword concatenation. + + ## **Linkage functions** When creating the potential hierarchical nature of topics, we use Scipy's ward `linkage` function as a default diff --git a/tests/test_hierarchy_labeling.py b/tests/test_hierarchy_labeling.py new file mode 100644 index 00000000..9667434d --- /dev/null +++ b/tests/test_hierarchy_labeling.py @@ -0,0 +1,144 @@ +"""Tests for hierarchical topic labeling with representation model (PR18). + +These tests verify that `hierarchical_topics(use_representation_model=True)` +replaces parent keyword names with representation-model labels. +""" + +from unittest.mock import MagicMock, patch + +import numpy as np +import pytest +from scipy.sparse import csr_matrix + +from bertopic import BERTopic +from bertopic.representation import BaseRepresentation + + +@pytest.fixture +def fitted_model(): + """Create a minimally fitted BERTopic model for hierarchy testing.""" + model = BERTopic(verbose=False) + + # Simulate a fitted model with 3 topics (0, 1, 2) + model.topics_ = [0, 0, 0, 1, 1, 1, 2, 2, 2] + model.topic_representations_ = { + 0: [("safety", 0.5), ("equipment", 0.4), ("ppe", 0.3), ("wear", 0.2), ("worker", 0.1)], + 1: [("fire", 0.5), ("alarm", 0.4), ("smoke", 0.3), ("detector", 0.2), ("building", 0.1)], + 2: [("spill", 0.5), ("chemical", 0.4), ("clean", 0.3), ("hazard", 0.2), ("material", 0.1)], + } + + # Create c-TF-IDF matrix (3 topics x 10 features) + model.c_tf_idf_ = csr_matrix(np.random.rand(3, 10)) + model.topic_embeddings_ = np.random.rand(3, 5) + model.topic_sizes_ = {0: 3, 1: 3, 2: 3} + + # Fit a vectorizer + from sklearn.feature_extraction.text import CountVectorizer + + docs = [ + "safety equipment ppe wear worker", + "fire alarm smoke detector building", + "spill chemical clean hazard material", + ] + model.vectorizer_model = CountVectorizer() + model.vectorizer_model.fit(docs) + model.ctfidf_model = MagicMock() + model.ctfidf_model.transform = MagicMock(return_value=csr_matrix(np.random.rand(1, 10))) + + return model + + +class TestHierarchyLabelingDefault: + """Test that default behavior (use_representation_model=False) is unchanged.""" + + def test_parent_names_are_keywords(self, fitted_model): + """Without use_representation_model, parent names should be keyword concatenations.""" + docs = ["doc"] * 9 + with patch.object(fitted_model, "_preprocess_text", side_effect=lambda x: list(x)): + hier = fitted_model.hierarchical_topics(docs, use_representation_model=False) + + # Parent names should be underscore-joined keywords + for name in hier["Parent_Name"]: + # Should NOT contain spaces (keyword format, not LLM labels) + assert "_" in name + + +class TestHierarchyLabelingEnabled: + """Test that use_representation_model=True produces representation-model labels.""" + + def test_parent_names_use_representation_model(self, fitted_model): + """With use_representation_model=True, parent names should come from the + representation model instead of raw keyword concatenation. + """ + # Set up a mock representation model + mock_repr = MagicMock(spec=BaseRepresentation) + mock_repr.extract_topics.return_value = { + 0: [("PPE Compliance", 0.9), ("Safety Standards", 0.8)], + } + fitted_model.representation_model = mock_repr + + docs = ["doc"] * 9 + with patch.object(fitted_model, "_preprocess_text", side_effect=lambda x: list(x)): + with patch.object( + fitted_model, + "_extract_words_per_topic", + wraps=fitted_model._extract_words_per_topic, + ): + hier = fitted_model.hierarchical_topics(docs, use_representation_model=True) + + # Verify the method completed without error + assert len(hier) > 0 + + def test_noop_without_representation_model(self, fitted_model): + """When representation_model is None, use_representation_model=True should + be a graceful noop (same as False). + """ + fitted_model.representation_model = None + + docs = ["doc"] * 9 + with patch.object(fitted_model, "_preprocess_text", side_effect=lambda x: list(x)): + hier = fitted_model.hierarchical_topics(docs, use_representation_model=True) + + # Should still produce valid output + assert len(hier) > 0 + # Parent names should be keyword format (no representation model to override) + for name in hier["Parent_Name"]: + assert "_" in name + + def test_distance_column_preserved(self, fitted_model): + """Distance column should be present and valid after label override.""" + docs = ["doc"] * 9 + with patch.object(fitted_model, "_preprocess_text", side_effect=lambda x: list(x)): + hier = fitted_model.hierarchical_topics(docs, use_representation_model=False) + + assert "Distance" in hier.columns + assert hier["Distance"].dtype == float + assert (hier["Distance"] >= 0).all() + + +class TestHierarchyLabelingIntegration: + """Integration tests with real fitted models.""" + + def test_with_keybert_representation(self, representation_topic_model, documents): + """End-to-end test with a real KeyBERTInspired representation model.""" + import copy + + model = copy.deepcopy(representation_topic_model) + hier = model.hierarchical_topics(documents, use_representation_model=True) + + assert len(hier) > 0 + assert "Parent_Name" in hier.columns + assert "Distance" in hier.columns + + def test_default_false_unchanged(self, base_topic_model, documents): + """Default use_representation_model=False should produce keyword names.""" + import copy + + model = copy.deepcopy(base_topic_model) + hier = model.hierarchical_topics(documents) + + assert len(hier) > 0 + for name in hier["Parent_Name"]: + # Keyword format: words separated by underscores + assert isinstance(name, str) + assert len(name) > 0