diff --git a/bertopic/_bertopic.py b/bertopic/_bertopic.py index cfafb58a..e9171032 100644 --- a/bertopic/_bertopic.py +++ b/bertopic/_bertopic.py @@ -1,7 +1,8 @@ # ruff: noqa: E402 -import yaml import warnings +import yaml + warnings.filterwarnings("ignore", category=FutureWarning) warnings.filterwarnings("ignore", category=UserWarning) @@ -10,26 +11,25 @@ except (KeyError, AttributeError, TypeError): pass -import re +import collections +import inspect import math +import re +from collections import Counter, defaultdict +from copy import deepcopy +from importlib.util import find_spec +from pathlib import Path +from tempfile import TemporaryDirectory +from typing import TYPE_CHECKING, Any, Callable, Iterable, List, Literal, Mapping, Tuple, Union + import joblib -import inspect -import collections import numpy as np import pandas as pd import scipy.sparse as sp -from copy import deepcopy - -from tqdm import tqdm -from pathlib import Path from packaging import version -from tempfile import TemporaryDirectory -from collections import defaultdict, Counter -from scipy.sparse import csr_matrix from scipy.cluster import hierarchy as sch -from importlib.util import find_spec - -from typing import List, Tuple, Union, Mapping, Any, Callable, Iterable, TYPE_CHECKING, Literal +from scipy.sparse import csr_matrix +from tqdm import tqdm # Plotting if find_spec("plotly") is None: @@ -41,8 +41,8 @@ from bertopic import plotting if TYPE_CHECKING: - import plotly.graph_objs as go import matplotlib.figure as fig + import plotly.graph_objs as go # Models @@ -54,32 +54,33 @@ HAS_HDBSCAN = False from sklearn.cluster import HDBSCAN as SK_HDBSCAN -from sklearn.preprocessing import normalize from sklearn import __version__ as sklearn_version from sklearn.cluster import AgglomerativeClustering from sklearn.decomposition import PCA -from sklearn.metrics.pairwise import cosine_similarity from sklearn.feature_extraction.text import CountVectorizer, TfidfTransformer +from sklearn.metrics.pairwise import cosine_similarity +from sklearn.preprocessing import normalize -# BERTopic -from bertopic.cluster import BaseCluster -from bertopic.backend import BaseEmbedder -from bertopic.representation._mmr import mmr -from bertopic.backend._utils import select_backend -from bertopic.vectorizers import ClassTfidfTransformer -from bertopic.representation import BaseRepresentation, KeyBERTInspired -from bertopic.dimensionality import BaseDimensionalityReduction -from bertopic.cluster._utils import hdbscan_delegator, is_supported_hdbscan +import bertopic._save_utils as save_utils from bertopic._utils import ( MyLogger, check_documents_type, check_embeddings_shape, check_is_fitted, - validate_distance_matrix, - select_topic_representation, get_unique_distances, + select_topic_representation, + validate_distance_matrix, ) -import bertopic._save_utils as save_utils +from bertopic.backend import BaseEmbedder +from bertopic.backend._utils import select_backend + +# BERTopic +from bertopic.cluster import BaseCluster +from bertopic.cluster._utils import hdbscan_delegator, is_supported_hdbscan +from bertopic.dimensionality import BaseDimensionalityReduction +from bertopic.representation import BaseRepresentation, KeyBERTInspired +from bertopic.representation._mmr import mmr +from bertopic.vectorizers import ClassTfidfTransformer logger = MyLogger() logger.configure("WARNING") @@ -1038,6 +1039,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, + strategy: str = "agglomerative", ) -> pd.DataFrame: """Create a hierarchy of topics. @@ -1063,6 +1065,9 @@ 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. + strategy: The hierarchy construction strategy. Either `"agglomerative"` + (bottom-up merging, the default) or `"divisive"` (top-down + splitting). Returns: hierarchical_topics: A dataframe that contains a hierarchy of topics @@ -1090,6 +1095,10 @@ def hierarchical_topics( ``` """ check_documents_type(docs) + + if strategy == "divisive": + return self._divisive_hierarchical_topics(docs, use_ctfidf=use_ctfidf) + if distance_function is None: distance_function = lambda x: 1 - cosine_similarity(x) @@ -1201,6 +1210,194 @@ def hierarchical_topics( return hier_topics + def _divisive_hierarchical_topics( + self, + docs: List[str], + use_ctfidf: bool = True, + ) -> pd.DataFrame: + """Build a topic hierarchy using top-down recursive binary splitting. + + At each step the set of topics with the highest internal diversity is split + into two children using rank-2 NMF on the c-TF-IDF matrix of the group. + The process repeats until every group contains a single topic. + + Arguments: + docs: The documents used when calling ``fit`` or ``fit_transform``. + use_ctfidf: Whether to use c-TF-IDF embeddings for distance calculations. + If ``False``, the embedding model's representations are used. + + Returns: + A DataFrame with the same schema as the agglomerative variant: + ``Parent_ID``, ``Parent_Name``, ``Topics``, ``Child_Left_ID``, + ``Child_Left_Name``, ``Child_Right_ID``, ``Child_Right_Name``, ``Distance``. + """ + from sklearn.decomposition import NMF + + # Prepare embeddings and documents + embeddings = select_topic_representation(self.c_tf_idf_, self.topic_embeddings_, use_ctfidf)[0][ + self._outliers : + ] + documents = pd.DataFrame({"Document": docs, "ID": range(len(docs)), "Topic": self.topics_}) + documents_per_topic = documents.groupby(["Topic"], as_index=False).agg({"Document": " ".join}) + documents_per_topic = documents_per_topic.loc[documents_per_topic.Topic != -1, :] + clean_documents = self._preprocess_text(documents_per_topic.Document.values) + + try: + words = self.vectorizer_model.get_feature_names_out() + except AttributeError: + words = self.vectorizer_model.get_feature_names() + + bow = self.vectorizer_model.transform(clean_documents) + + # Topic list (excluding outliers) + topic_ids = sorted([t for t in set(self.topics_) if t != -1]) + n_topics = len(topic_ids) + + # Internal node IDs start at n_topics (leaves are 0..n_topics-1). + # We assign IDs incrementally, then remap so the root (first split) + # gets the highest ID — matching the agglomerative convention that + # get_topic_tree expects. + next_id = n_topics + + # Track each group's assigned internal ID + group_to_id = {} + initial_key = tuple(range(n_topics)) + group_to_id[initial_key] = next_id + next_id += 1 + + groups = [list(range(n_topics))] + rows = [] + + def _child_name(indices): + """Derive a short name for a child node.""" + if len(indices) == 1: + return "_".join([x[0] for x in self.get_topic(topic_ids[indices[0]])][:5]) + child_bow = csr_matrix(bow[indices].sum(axis=0)) + child_c_tf_idf = self.ctfidf_model.transform(child_bow) + child_sel = documents.loc[documents.Topic.isin([topic_ids[i] for i in indices]), :] + child_sel = child_sel.copy() + child_sel.Topic = 0 + child_wpt = self._extract_words_per_topic(words, child_sel, child_c_tf_idf, calculate_aspects=False) + return "_".join([x[0] for x in child_wpt[0]][:5]) + + # Iteratively split the most diverse group + while any(len(g) > 1 for g in groups): + # Find the group with the highest internal diversity + best_group_idx = None + best_diversity = -1 + for gi, group in enumerate(groups): + if len(group) <= 1: + continue + group_emb = embeddings[group] + sim = cosine_similarity(group_emb) + diversity = 1 - sim[np.triu_indices_from(sim, k=1)].mean() + if diversity > best_diversity: + best_diversity = diversity + best_group_idx = gi + + if best_group_idx is None: + break + + group = groups.pop(best_group_idx) + parent_id = group_to_id[tuple(sorted(group))] + + # Split using rank-2 NMF on the group's c-TF-IDF + group_c_tf_idf = self.ctfidf_model.transform(csr_matrix(np.vstack([bow[idx].toarray() for idx in group]))) + + if group_c_tf_idf.shape[0] < 2: + groups.append(group) + continue + + try: + nmf = NMF(n_components=2, random_state=42, max_iter=300) + W = nmf.fit_transform(group_c_tf_idf) + except ValueError: + # NMF can fail on degenerate input (e.g., all-zero rows); + # fall back to a simple positional split. + W = None + + if W is not None: + # Assign each topic to the component with highest weight + assignments = np.argmax(W, axis=1) + left_indices = [group[i] for i in range(len(group)) if assignments[i] == 0] + right_indices = [group[i] for i in range(len(group)) if assignments[i] == 1] + + # Ensure non-empty splits + if not left_indices or not right_indices: + mid = len(group) // 2 + left_indices, right_indices = group[:mid], group[mid:] + else: + mid = len(group) // 2 + left_indices, right_indices = group[:mid], group[mid:] + + # Parent name from c-TF-IDF of combined group + grouped_bow = csr_matrix(bow[group].sum(axis=0)) + c_tf_idf = self.ctfidf_model.transform(grouped_bow) + selection = documents.loc[documents.Topic.isin([topic_ids[i] for i in group]), :] + selection = selection.copy() + selection.Topic = 0 + words_per_topic = self._extract_words_per_topic(words, selection, c_tf_idf, calculate_aspects=False) + parent_name = "_".join([x[0] for x in words_per_topic[0]][:5]) + + # Assign child IDs: leaves get their topic ID, non-leaves get + # a fresh internal ID allocated immediately. + if len(left_indices) == 1: + left_id = topic_ids[left_indices[0]] + else: + left_id = next_id + group_to_id[tuple(sorted(left_indices))] = next_id + next_id += 1 + + if len(right_indices) == 1: + right_id = topic_ids[right_indices[0]] + else: + right_id = next_id + group_to_id[tuple(sorted(right_indices))] = next_id + next_id += 1 + + left_name = _child_name(left_indices) + right_name = _child_name(right_indices) + + rows.append( + { + "Parent_ID": parent_id, + "Parent_Name": parent_name, + "Topics": [topic_ids[i] for i in group], + "Child_Left_ID": left_id, + "Child_Left_Name": left_name, + "Child_Right_ID": right_id, + "Child_Right_Name": right_name, + "Distance": best_diversity, + } + ) + + # Add children back as groups + groups.append(left_indices) + groups.append(right_indices) + + # Remap internal node IDs so the root (first split) has the highest + # ID. get_topic_tree identifies the root via max(Parent_ID). + if rows: + id_min = n_topics + id_max = next_id - 1 + + def _remap(node_id): + return node_id if node_id < n_topics else id_min + id_max - node_id + + for row in rows: + row["Parent_ID"] = _remap(row["Parent_ID"]) + row["Child_Left_ID"] = _remap(row["Child_Left_ID"]) + row["Child_Right_ID"] = _remap(row["Child_Right_ID"]) + + hier_topics = pd.DataFrame(rows) + if not hier_topics.empty: + hier_topics = hier_topics.sort_values("Parent_ID", ascending=False) + hier_topics[["Parent_ID", "Child_Left_ID", "Child_Right_ID"]] = hier_topics[ + ["Parent_ID", "Child_Left_ID", "Child_Right_ID"] + ].astype(str) + + return hier_topics + def approximate_distribution( self, documents: Union[str, List[str]], diff --git a/docs/getting_started/hierarchicaltopics/hierarchicaltopics.md b/docs/getting_started/hierarchicaltopics/hierarchicaltopics.md index 8780b668..b6bfa82c 100644 --- a/docs/getting_started/hierarchicaltopics/hierarchicaltopics.md +++ b/docs/getting_started/hierarchicaltopics/hierarchicaltopics.md @@ -341,6 +341,34 @@ to view, we can see better which topics could be logically merged: +## **Divisive hierarchy** + +By default, `hierarchical_topics` uses an agglomerative (bottom-up) strategy. +You can also use a divisive (top-down) strategy that recursively splits topics: + +```python +hierarchical_topics = topic_model.hierarchical_topics(docs, strategy="divisive") +``` + +The divisive strategy produces a binary tree where each split maximizes the +distance between the resulting sub-clusters. This can be more intuitive when you +have a natural hierarchy of broad-to-specific topics. + +Both strategies produce the same DataFrame format, so `get_topic_tree()` and +`visualize_hierarchy()` work with either: + +```python +# Works with both strategies +tree = topic_model.get_topic_tree(hierarchical_topics) +fig = topic_model.visualize_hierarchy(hierarchical_topics=hierarchical_topics) +``` + +!!! tip + The agglomerative strategy tends to work better when topics are + well-separated, while the divisive strategy can be better at identifying + fine-grained sub-topics within broad categories. + + ## **Merge topics** After seeing the potential hierarchy of your topic, you might want to merge specific diff --git a/tests/test_divisive_hierarchy.py b/tests/test_divisive_hierarchy.py new file mode 100644 index 00000000..4cd1b230 --- /dev/null +++ b/tests/test_divisive_hierarchy.py @@ -0,0 +1,203 @@ +"""Tests for divisive hierarchical_topics strategy. + +Run from BERTopic repo root: + pytest tests/test_divisive_hierarchy.py -v +""" + +import copy +from unittest.mock import MagicMock + +import numpy as np +import pandas as pd +from scipy.sparse import csr_matrix + + +# --------------------------------------------------------------------------- +# Helper: build a minimal mock BERTopic model for unit tests +# --------------------------------------------------------------------------- + + +def _make_mock_model(n_topics, n_docs): + """Create a minimal mock BERTopic model for testing divisive hierarchy.""" + from bertopic._bertopic import BERTopic + + model = MagicMock(spec=BERTopic) + model._outliers = 1 + model.topics_ = [i % n_topics for i in range(n_docs)] + + # Topic embeddings: outlier + n_topics + np.random.seed(42) + model.topic_embeddings_ = np.random.rand(n_topics + 1, 10) + model.c_tf_idf_ = csr_matrix(np.random.rand(n_topics + 1, 20)) + + # Vectorizer mock + model.vectorizer_model = MagicMock() + model.vectorizer_model.get_feature_names_out.return_value = np.array([f"word{i}" for i in range(20)]) + model.vectorizer_model.transform.return_value = csr_matrix(np.random.rand(n_topics, 20)) + + # ctfidf model + model.ctfidf_model = MagicMock() + model.ctfidf_model.transform.side_effect = lambda x: csr_matrix(np.random.rand(x.shape[0], 20)) + + # Preprocessing + model._preprocess_text.side_effect = lambda x: x + + # get_topic returns word tuples + model.get_topic.return_value = [(f"word{i}", 0.5 - i * 0.1) for i in range(5)] + + # _extract_words_per_topic returns dict + model._extract_words_per_topic.return_value = {0: [(f"word{i}", 0.5 - i * 0.1) for i in range(5)]} + + # Bind the real method + model._divisive_hierarchical_topics = BERTopic._divisive_hierarchical_topics.__get__(model) + + return model + + +# --------------------------------------------------------------------------- +# Unit tests (mock model) +# --------------------------------------------------------------------------- + + +def test_divisive_returns_dataframe(): + """Divisive strategy should return a DataFrame with expected columns.""" + model = _make_mock_model(n_topics=4, n_docs=20) + result = model._divisive_hierarchical_topics(docs=["doc"] * 20, use_ctfidf=True) + assert isinstance(result, pd.DataFrame) + expected_cols = { + "Parent_ID", + "Parent_Name", + "Topics", + "Child_Left_ID", + "Child_Left_Name", + "Child_Right_ID", + "Child_Right_Name", + "Distance", + } + assert set(result.columns) == expected_cols + + +def test_divisive_two_topics(): + """With 2 topics, should produce exactly 1 split.""" + model = _make_mock_model(n_topics=2, n_docs=10) + result = model._divisive_hierarchical_topics(docs=["doc"] * 10, use_ctfidf=True) + assert len(result) == 1 + + +def test_divisive_single_topic_empty(): + """With 1 topic, no splits needed — empty or no rows.""" + model = _make_mock_model(n_topics=1, n_docs=5) + result = model._divisive_hierarchical_topics(docs=["doc"] * 5, use_ctfidf=True) + assert len(result) == 0 + + +def test_divisive_ids_are_strings(): + """Parent and child IDs should be string type.""" + model = _make_mock_model(n_topics=3, n_docs=15) + result = model._divisive_hierarchical_topics(docs=["doc"] * 15, use_ctfidf=True) + if not result.empty: + assert result["Parent_ID"].dtype == object # string + assert result["Child_Left_ID"].dtype == object + assert result["Child_Right_ID"].dtype == object + + +def test_divisive_distances_non_negative(): + """All distances should be non-negative.""" + model = _make_mock_model(n_topics=5, n_docs=25) + result = model._divisive_hierarchical_topics(docs=["doc"] * 25, use_ctfidf=True) + if not result.empty: + assert (result["Distance"] >= 0).all() + + +def test_divisive_no_id_collisions(): + """Internal node IDs must be unique — no collisions on depth > 2 trees.""" + model = _make_mock_model(n_topics=8, n_docs=40) + result = model._divisive_hierarchical_topics(docs=["doc"] * 40, use_ctfidf=True) + if result.empty: + return + + # Collect all node IDs that appear in the hierarchy + all_ids = set() + for _, row in result.iterrows(): + for col in ("Parent_ID", "Child_Left_ID", "Child_Right_ID"): + all_ids.add(row[col]) + + # Parent IDs must all be unique + parent_ids = result["Parent_ID"].tolist() + assert len(parent_ids) == len(set(parent_ids)), "Duplicate Parent_IDs found" + + +def test_divisive_root_has_max_parent_id(): + """The root node should have the highest Parent_ID (get_topic_tree convention).""" + model = _make_mock_model(n_topics=6, n_docs=30) + result = model._divisive_hierarchical_topics(docs=["doc"] * 30, use_ctfidf=True) + if result.empty: + return + + # The root split should contain all topics + n_topics = len([t for t in set(model.topics_) if t != -1]) + root_row = result.loc[result["Parent_ID"].astype(int).idxmax()] + assert len(root_row["Topics"]) == n_topics + + +def test_divisive_leaf_ids_below_n_topics(): + """Leaf (single-topic) child IDs should be actual topic IDs, not internal IDs.""" + model = _make_mock_model(n_topics=4, n_docs=20) + result = model._divisive_hierarchical_topics(docs=["doc"] * 20, use_ctfidf=True) + if result.empty: + return + + topic_ids = sorted([t for t in set(model.topics_) if t != -1]) + + # Gather IDs that are NOT parent IDs (i.e., they are leaves) + parent_ids = set(result["Parent_ID"].tolist()) + for _, row in result.iterrows(): + for col in ("Child_Left_ID", "Child_Right_ID"): + child_id = row[col] + if child_id not in parent_ids: + # This is a leaf — its int ID should be a valid topic ID + assert int(child_id) in topic_ids + + +# --------------------------------------------------------------------------- +# Integration tests (real fitted model from conftest fixtures) +# --------------------------------------------------------------------------- + + +def test_divisive_strategy_on_fitted_model(base_topic_model, documents): + """hierarchical_topics(strategy='divisive') should work on a fitted model.""" + model = copy.deepcopy(base_topic_model) + result = model.hierarchical_topics(documents, strategy="divisive") + + assert isinstance(result, pd.DataFrame) + expected_cols = { + "Parent_ID", + "Parent_Name", + "Topics", + "Child_Left_ID", + "Child_Left_Name", + "Child_Right_ID", + "Child_Right_Name", + "Distance", + } + assert set(result.columns) == expected_cols + assert len(result) > 0 + + +def test_default_strategy_is_agglomerative(base_topic_model, documents): + """Default strategy should be agglomerative (backward compat).""" + model = copy.deepcopy(base_topic_model) + result_default = model.hierarchical_topics(documents) + result_agg = model.hierarchical_topics(documents, strategy="agglomerative") + + # Same columns and structure + assert list(result_default.columns) == list(result_agg.columns) + + +def test_get_topic_tree_compat(base_topic_model, documents): + """get_topic_tree should work with divisive hierarchy output.""" + model = copy.deepcopy(base_topic_model) + hier = model.hierarchical_topics(documents, strategy="divisive") + tree = model.get_topic_tree(hier) + assert isinstance(tree, str) + assert len(tree) > 0