diff --git a/bertopic/_bertopic.py b/bertopic/_bertopic.py index cfafb58a..36ad4827 100644 --- a/bertopic/_bertopic.py +++ b/bertopic/_bertopic.py @@ -1201,6 +1201,362 @@ def hierarchical_topics( return hier_topics + def create_topic_taxonomy( + self, + docs: List[str], + nr_topics_per_level: List[int], + embeddings: np.ndarray = None, + min_children: int = 2, + doc_embedding_weight: float = 0.5, + use_ctfidf: bool = False, + ) -> pd.DataFrame: + """Builds a deterministic N-level taxonomy on top of the fitted leaf topics. + Each level is created by clustering the level below using agglomerative + clustering with cosine distance, then repairing any clusters that have + fewer than ``min_children`` by merging them into the most similar neighbor. + + The leaf-level topic model remains unchanged - parent topics are derived + through a lookup table and are not used during inference. + + Arguments: + docs: The documents you used when calling either ``fit`` or ``fit_transform`` + nr_topics_per_level: Target number of topics for each level above the leaves. + For example, ``[20, 5]`` creates two parent levels: level 1 with ~20 topics + and level 2 with ~5 topics. Each entry must be less than the number of + topics at the level below. + embeddings: Pre-trained document embeddings used to compute document centroids + per topic. If ``None``, only BERTopic's internal topic embeddings are used + and ``doc_embedding_weight`` is ignored. + min_children: Minimum number of children each parent must have at every level. + Clusters with fewer children are merged into their most similar neighbor. + Must be >= 1. Default is 2. + doc_embedding_weight: Weight for document centroids when computing leaf topic + vectors. The topic embedding weight is ``1 - doc_embedding_weight``. + Only used when ``embeddings`` is provided. Must be in [0, 1]. Default is 0.5. + use_ctfidf: Whether to use c-TF-IDF representations for clustering distance + computation. If ``False``, topic embedding vectors are used. Default is ``False``. + + Returns: + hierarchy: A DataFrame with the full hierarchy. Columns: + + - ``Topic_ID``: Unique topic identifier. Leaf IDs match the fitted model's + topic IDs. Parent IDs are auto-incremented integers. + - ``Topic_Name``: Top-5 c-TF-IDF words joined by ``_``. + - ``Level``: Hierarchy level (0 = leaf, 1 = first parent, etc.). + - ``Parent_ID``: ID of the parent topic at the next level up. ``-2`` for + top-level topics. + - ``Parent_Name``: Name of the parent topic. ``"Root"`` for top-level. + - ``Document_Count``: Number of documents. For parents this is the sum of + their children's counts. + + Examples: + ```python + from bertopic import BERTopic + topic_model = BERTopic() + topics, probs = topic_model.fit_transform(docs) + + # Two-level hierarchy: ~20 parents, ~5 grandparents + hierarchy = topic_model.create_topic_taxonomy( + docs, nr_topics_per_level=[20, 5], min_children=2 + ) + + # With document embeddings for richer topic vectors + hierarchy = topic_model.create_topic_taxonomy( + docs, nr_topics_per_level=[20, 5], + embeddings=embeddings, doc_embedding_weight=0.7 + ) + ``` + """ + check_is_fitted(self) + check_documents_type(docs) + + # Validate parameters + if not isinstance(nr_topics_per_level, list) or len(nr_topics_per_level) == 0: + raise ValueError("`nr_topics_per_level` must be a non-empty list of positive integers.") + if any(n < 1 for n in nr_topics_per_level): + raise ValueError("All entries in `nr_topics_per_level` must be >= 1.") + if min_children < 1: + raise ValueError("`min_children` must be >= 1.") + if not 0.0 <= doc_embedding_weight <= 1.0: + raise ValueError("`doc_embedding_weight` must be between 0.0 and 1.0.") + + # Non-outlier leaf topics + topic_ids = sorted([t for t in self.topic_sizes_.keys() if t != -1]) + n_leaves = len(topic_ids) + if n_leaves < 2: + raise ValueError("Cannot create hierarchy with fewer than 2 non-outlier topics.") + + # Validate level targets against available topics + n_at_level = n_leaves + for k, target_n in enumerate(nr_topics_per_level): + if target_n >= n_at_level: + raise ValueError( + f"nr_topics_per_level[{k}] = {target_n} must be less than the " + f"number of topics at the level below ({n_at_level})." + ) + n_at_level = target_n + + if embeddings is not None: + check_embeddings_shape(embeddings, docs) + + # Phase 1: Prepare leaf level + topic_to_idx = {tid: i for i, tid in enumerate(topic_ids)} + topic_embeds = self.topic_embeddings_[self._outliers :] + + # Compute leaf topic vectors + if embeddings is not None: + doc_topics = np.array(self.topics_) + doc_centroids = np.zeros_like(topic_embeds) + for tid in topic_ids: + mask = doc_topics == tid + doc_centroids[topic_to_idx[tid]] = embeddings[mask].mean(axis=0) + w = doc_embedding_weight + level_vectors = w * doc_centroids + (1.0 - w) * topic_embeds + else: + level_vectors = topic_embeds.copy() + + doc_counts = np.array([self.topic_sizes_[tid] for tid in topic_ids]) + + # Prepare BoW for label generation + 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) + + # Support older sklearn versions, as was done in other parts as well. + if version.parse(sklearn_version) >= version.parse("1.0.0"): + words = self.vectorizer_model.get_feature_names_out() + else: + words = self.vectorizer_model.get_feature_names() + + bow = self.vectorizer_model.transform(clean_documents) + + # Reorder BoW rows to match topic_ids order + bow_topic_order = documents_per_topic.Topic.tolist() + bow_idx_map = {tid: i for i, tid in enumerate(bow_topic_order)} + reorder = [bow_idx_map[tid] for tid in topic_ids] + bow = bow[reorder] + + # Initialize hierarchy records with leaf topics + has_outliers = self._outliers == 1 + records = [] + + # Add outlier topic as a leaf if it exists + if has_outliers: + outlier_label = self.topic_labels_.get(-1, "-1_Outlier") + outlier_doc_count = self.topic_sizes_.get(-1, 0) + records.append( + { + "Topic_ID": -1, + "Topic_Name": outlier_label, + "Level": 0, + "Parent_ID": None, + "Parent_Name": None, + "Document_Count": outlier_doc_count, + } + ) + + for tid in topic_ids: + label = self.topic_labels_.get(tid, f"Topic_{tid}") + records.append( + { + "Topic_ID": tid, + "Topic_Name": label, + "Level": 0, + "Parent_ID": None, + "Parent_Name": None, + "Document_Count": self.topic_sizes_[tid], + } + ) + + # Track current level state (outliers are handled separately) + current_ids = list(topic_ids) + current_vectors = level_vectors + current_doc_counts = doc_counts + current_bow = bow + next_parent_id = max(topic_ids) + 1 + # Track which original leaf topic IDs each node covers (for BoW/doc lookups) + current_leaf_sets = [[tid] for tid in topic_ids] + # Track the outlier node ID at each level so we can chain it upward + outlier_child_id = -1 if has_outliers else None + + # Phase 2: Build each parent level + for k, target_n in enumerate(nr_topics_per_level, start=1): + n_current = len(current_ids) + + # Clamp target to actual available topics (prior-level repair may have reduced the count) + if target_n >= n_current: + target_n = n_current - 1 + if target_n < 1: + break + + # Compute distance matrix + if use_ctfidf: + c_tf_idf_level = self.ctfidf_model.transform(current_bow) + if isinstance(c_tf_idf_level, csr_matrix): + c_tf_idf_level = c_tf_idf_level.toarray() + dist_matrix = 1 - cosine_similarity(c_tf_idf_level) + else: + dist_matrix = 1 - cosine_similarity(current_vectors) + np.fill_diagonal(dist_matrix, 0) + + # Agglomerative clustering (again, support older versions) + if version.parse(sklearn_version) >= version.parse("1.4.0"): + agg = AgglomerativeClustering(target_n, metric="precomputed", linkage="average") + else: + agg = AgglomerativeClustering(target_n, affinity="precomputed", linkage="average") + agg.fit(dist_matrix) + + # Build cluster-to-children mapping + cluster_to_children = defaultdict(list) + for i, cl in enumerate(agg.labels_): + cluster_to_children[cl].append(i) + + # "Repair": merge clusters with fewer than min_children + changed = True + while changed: + changed = False + centroids = {} + for cl, children_idxs in cluster_to_children.items(): + weights = current_doc_counts[children_idxs] + centroids[cl] = np.average(current_vectors[children_idxs], axis=0, weights=weights) + + small_clusters = [cl for cl, ch in cluster_to_children.items() if len(ch) < min_children] + for cl in small_clusters: + if cl not in cluster_to_children or len(cluster_to_children) <= 1: + continue + + # Find most similar neighbor by cosine similarity + cl_centroid = centroids[cl].reshape(1, -1) + best_neighbor = None + best_sim = -1.0 + for other_cl, other_centroid in centroids.items(): + if other_cl == cl: + continue + sim = cosine_similarity(cl_centroid, other_centroid.reshape(1, -1))[0, 0] + if sim > best_sim: + best_sim = sim + best_neighbor = other_cl + + # Merge into best neighbor + cluster_to_children[best_neighbor].extend(cluster_to_children[cl]) + del cluster_to_children[cl] + # Recompute merged centroid + idxs = cluster_to_children[best_neighbor] + weights = current_doc_counts[idxs] + centroids[best_neighbor] = np.average(current_vectors[idxs], axis=0, weights=weights) + del centroids[cl] + changed = True + + # Build parent nodes for this level + next_level_ids = [] + next_level_vectors = [] + next_level_doc_counts = [] + next_level_bow = [] + next_level_leaf_sets = [] + + for cl in sorted(cluster_to_children.keys()): + children_idxs = cluster_to_children[cl] + parent_id = next_parent_id + next_parent_id += 1 + + # Collect all original leaf topic IDs covered by this parent + parent_leaf_ids = [] + for i in children_idxs: + parent_leaf_ids.extend(current_leaf_sets[i]) + + # Compute parent label via c-TF-IDF on merged BoW + parent_bow = csr_matrix(current_bow[children_idxs].sum(axis=0)) + c_tf_idf_parent = self.ctfidf_model.transform(parent_bow) + temp_docs = documents.loc[documents.Topic.isin(parent_leaf_ids)].copy() + temp_docs["Topic"] = 0 + words_per_topic = self._extract_words_per_topic( + words, + temp_docs, + c_tf_idf_parent, + calculate_aspects=False, + fine_tune_representation=False, + ) + parent_name = "_".join([x[0] for x in words_per_topic[0]][:5]) + + # Compute parent vector and doc count + weights = current_doc_counts[children_idxs] + parent_vector = np.average(current_vectors[children_idxs], axis=0, weights=weights) + parent_doc_count = int(current_doc_counts[children_idxs].sum()) + + # Update children's parent references + for idx in children_idxs: + child_id = current_ids[idx] + for rec in records: + if rec["Topic_ID"] == child_id and rec["Level"] == k - 1: + rec["Parent_ID"] = parent_id + rec["Parent_Name"] = parent_name + break + + records.append( + { + "Topic_ID": parent_id, + "Topic_Name": parent_name, + "Level": k, + "Parent_ID": None, + "Parent_Name": None, + "Document_Count": parent_doc_count, + } + ) + + next_level_ids.append(parent_id) + next_level_vectors.append(parent_vector) + next_level_doc_counts.append(parent_doc_count) + next_level_bow.append(parent_bow) + next_level_leaf_sets.append(parent_leaf_ids) + + # Create outlier parent for this level (single-child chain, bypasses clustering) + if has_outliers and outlier_child_id is not None: + outlier_parent_id = next_parent_id + next_parent_id += 1 + outlier_name = "Outlier" + + # Link the outlier child to this outlier parent + for rec in records: + if rec["Topic_ID"] == outlier_child_id and rec["Parent_ID"] is None: + rec["Parent_ID"] = outlier_parent_id + rec["Parent_Name"] = outlier_name + break + + records.append( + { + "Topic_ID": outlier_parent_id, + "Topic_Name": outlier_name, + "Level": k, + "Parent_ID": None, + "Parent_Name": None, + "Document_Count": outlier_doc_count, + } + ) + outlier_child_id = outlier_parent_id + + # Prepare for next level + current_ids = next_level_ids + current_vectors = np.array(next_level_vectors) + current_doc_counts = np.array(next_level_doc_counts) + current_bow = sp.vstack(next_level_bow) + current_leaf_sets = next_level_leaf_sets + + # Stop if we can't cluster further + if len(current_ids) <= 1: + break + + # Phase 3: Finalize top-level parents (including outlier chain) + max_level = max(rec["Level"] for rec in records) + for rec in records: + if rec["Level"] == max_level and rec["Parent_ID"] is None: + rec["Parent_ID"] = -2 + rec["Parent_Name"] = "Root" + + hierarchy_df = pd.DataFrame(records) + hierarchy_df = hierarchy_df.sort_values(["Level", "Topic_ID"]).reset_index(drop=True) + return hierarchy_df + def approximate_distribution( self, documents: Union[str, List[str]], diff --git a/docs/getting_started/taxonomy/taxonomy.md b/docs/getting_started/taxonomy/taxonomy.md new file mode 100644 index 00000000..5ff91be3 --- /dev/null +++ b/docs/getting_started/taxonomy/taxonomy.md @@ -0,0 +1,188 @@ +While `.hierarchical_topics()` explores topic similarity through a dendrogram and `.reduce_topics()` merges topics to a target count, `.create_topic_taxonomy()` takes a different approach: it builds a structured, multi-level taxonomy where you control the number of levels, the number of topics per level, and the minimum number of children each parent must have. + +The leaf topics produced by BERTopic stay untouched and are still used for inference. Parent topics are constructed bottom-up through agglomerative clustering, one level at a time, producing a lookup table that maps every leaf to its parent, grandparent, and so on. + + +## **Example** + +We train a BERTopic model on the 20 Newsgroups dataset and then build a two-level taxonomy on top of the resulting leaf topics: + +```python +from bertopic import BERTopic +from sklearn.datasets import fetch_20newsgroups + +docs = fetch_20newsgroups(subset='all', remove=('headers', 'footers', 'quotes'))["data"] + +topic_model = BERTopic(verbose=True) +topics, probs = topic_model.fit_transform(docs) + +hierarchy = topic_model.create_topic_taxonomy( + docs, + nr_topics_per_level=[20, 5], + min_children=2, +) +``` + +The result is a DataFrame describing the full taxonomy. Each row contains a topic's ID, name, hierarchy level, and a reference to its parent. Top-level topics have `Parent_ID = -2`: + +```python +>>> print(hierarchy.to_string(index=False)) + + Topic_ID Topic_Name Level Parent_ID Parent_Name Document_Count + 0 0_game_team_games_he 0 241 game_team_he_games_players 1834 + 1 1_key_clipper_chip_encryption 0 253 key_encryption_clipper_chip_keys 603 + 2 2_israel_israeli_jews_arab 0 240 armenian_were_armenians_turkish_they 459 + 8 8_car_cars_mustang_ford 0 265 car_bike_cars_engine_miles 169 + 13 13_bike_bikes_miles_honda 0 265 car_bike_cars_engine_miles 169 + ... + 265 car_bike_cars_engine_miles 1 276 bike_car_insurance_my_you 1104 + ... + 276 bike_car_insurance_my_you 2 -2 Root 3820 + ... +``` + +Here, leaf topics 8 and 13 both map to parent 265 (`car_bike_cars_engine_miles`), and parent 265 maps to grandparent 276 (`bike_car_insurance_my_you`). You can verify the structure: + +```python +>>> hierarchy.groupby("Level").size() +Level +0 225 +1 45 +2 15 +dtype: int64 +``` + + +## **How it works** + +### 1. Compute leaf topic vectors + +For every leaf topic (so the ones generated by BERTopic), a topic vector is computed. **If document embeddings are provided**, the vector is a weighted blend of the document centroid and BERTopic's internal topic embedding: + +``` +Topic Vector = w * Document Centroid + (1 - w) * Topic Embedding +``` + +The weight `w` is controlled by `doc_embedding_weight` (default 0.5). Without document embeddings, the topic embedding is used directly. It's recommended to at least experiment with using the document embeddings, as the produced taxonomy tends to be more accurate through it. + +### 2. Build each parent level + +For each level above the leaves: + +1. **Cluster** the topic vectors using agglomerative clustering with cosine distance +2. **Repair** any cluster with fewer than `min_children` topics by merging it into the most similar neighbor. Centroids are recomputed after each merge. +3. **Label** each parent using c-TF-IDF on the combined bag-of-words of its children +4. **Compute** the parent vector as a weighted centroid of its children (weighted by document count) + +### 3. Finalize + +Top-level topics are marked with `Parent_ID = -2`. + + +## **Parameters** + +```python +hierarchy = topic_model.create_topic_taxonomy( + docs, + nr_topics_per_level=[20, 5], + embeddings=embeddings, # optional document embeddings + min_children=2, # minimum children per parent + doc_embedding_weight=0.5, # blend weight for doc centroids + use_ctfidf=False, # use c-TF-IDF for distance computation +) +``` + +| Parameter | Default | Description | +|---|---|-----------------------------------------------------------------------------------------------------------------------------------------------------------| +| `nr_topics_per_level` | *required* | Target topic count for each level above the leaves. E.g., `[20, 5]` creates two parent levels. Each entry must be less than the count at the level below. | +| `embeddings` | `None` | Document embeddings for computing per-topic centroids (recommended). If not provided, only BERTopic's topic embeddings are used. | +| `min_children` | `2` | Minimum children per parent at every level. Clusters with fewer children are merged into their most similar neighbor. | +| `doc_embedding_weight` | `0.5` | Weight for document centroids vs. topic embeddings. Only used when `embeddings` is provided. | +| `use_ctfidf` | `False` | Whether to use c-TF-IDF representations instead of topic embeddings for distance computation. | + + +## **Why the topic count is approximate** + +The values in `nr_topics_per_level` are targets, not exact guarantees. The actual number of topics at each level may end up slightly lower. This is because of the repair step. + +The algorithm first creates exactly the requested number of clusters using agglomerative clustering. But some of those clusters may contain fewer than `min_children` topics. The repair step then merges each undersized cluster into its most similar neighbor, which reduces the total count. + +For example, if you request 20 parents from 60 leaves with `min_children=2`, the clustering produces 20 clusters. If three of those clusters happen to contain only a single leaf, they get merged into nearby clusters, leaving you with 17 parents instead of 20. + +This effect becomes more pronounced as you increase `nr_topics_per_level` relative to the number of topics at the level below. More clusters means fewer topics per cluster on average, which means more clusters fall below the `min_children` threshold and get merged away. In practice, keeping the ratio between consecutive levels at roughly 3:1 or higher tends to produce counts close to the target. + +!!! note + If you set `min_children=1`, no repair merging takes place and you will get exactly the requested number of topics at each level. + + +## **Using document embeddings** + +If you saved the document embeddings during training, you can pass them in for richer topic vectors: + +```python +embedding_model = SentenceTransformer("all-MiniLM-L6-v2") +embeddings = embedding_model.encode(docs) + +topic_model = BERTopic(embedding_model=embedding_model) +topics, probs = topic_model.fit_transform(docs, embeddings) + +hierarchy = topic_model.create_topic_taxonomy( + docs, + nr_topics_per_level=[20, 5], + embeddings=embeddings, + doc_embedding_weight=0.7, +) +``` + +A higher `doc_embedding_weight` puts more emphasis on how documents are distributed in embedding space. A lower weight relies more on BERTopic's topic embeddings, which capture word-level topic structure. + + +## **Controlling taxonomy shape** + +The `min_children` parameter ensures every parent has enough children to be meaningful: + +``` +min_children=2 (default) min_children=3 + +Grandparent A Grandparent A + ├── Parent 1 ├── Parent 1 + │ ├── Leaf 1 │ ├── Leaf 1 + │ └── Leaf 2 │ ├── Leaf 2 + └── Parent 2 │ └── Leaf 3 + ├── Leaf 3 └── Parent 2 + └── Leaf 4 ├── Leaf 4 + ├── Leaf 5 + └── Leaf 6 +``` + +When a cluster has fewer children than required, it is merged into its most similar neighbor based on cosine similarity between their centroids, iteratively, until all clusters satisfy the constraint. + +```python +hierarchy = topic_model.create_topic_taxonomy( + docs, + nr_topics_per_level=[15, 5], + min_children=3, +) +``` + +!!! tip + `create_topic_taxonomy()` does **not** modify the fitted model. Your leaf topics, representations, and probabilities remain exactly as they were. The taxonomy is a read-only lookup table on top of the existing model. + + +## **Navigating the taxonomy** + +The returned DataFrame makes it straightforward to traverse the hierarchy: + +```python +# Get all parents at level 1 +parents = hierarchy[hierarchy.Level == 1] + +# Find which leaves belong to a specific parent +parent_id = parents.iloc[0].Topic_ID +children = hierarchy[hierarchy.Parent_ID == parent_id] + +# Trace a leaf topic up to the root +leaf = hierarchy[hierarchy.Topic_ID == 0] +parent = hierarchy[hierarchy.Topic_ID == leaf.iloc[0].Parent_ID] +grandparent = hierarchy[hierarchy.Topic_ID == parent.iloc[0].Parent_ID] +``` diff --git a/mkdocs.yml b/mkdocs.yml index 26a823df..4e08073d 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -40,6 +40,7 @@ nav: - Variations: - Dynamic Topic Modeling: getting_started/topicsovertime/topicsovertime.md - Hierarchical Topic Modeling: getting_started/hierarchicaltopics/hierarchicaltopics.md + - Topic Taxonomy: getting_started/taxonomy/taxonomy.md - Multimodal Topic Modeling: getting_started/multimodal/multimodal.md - Online Topic Modeling: getting_started/online/online.md - Merge Multiple Models: getting_started/merge/merge.md diff --git a/tests/test_variations/test_taxonomy.py b/tests/test_variations/test_taxonomy.py new file mode 100644 index 00000000..39b96337 --- /dev/null +++ b/tests/test_variations/test_taxonomy.py @@ -0,0 +1,278 @@ +import copy +import pytest +import numpy as np +import pandas as pd + + +@pytest.mark.parametrize( + "model", + [ + "base_topic_model", + "custom_topic_model", + "reduced_topic_model", + ], +) +class TestCreateTopicTaxonomy: + """Tests for the create_topic_taxonomy method.""" + + def _get_model_and_nr_topics(self, model, request): + """Get a deep copy of the model and count non-outlier topics.""" + topic_model = copy.deepcopy(request.getfixturevalue(model)) + nr_topics = len([t for t in topic_model.topic_sizes_.keys() if t != -1]) + return topic_model, nr_topics + + def test_basic_two_levels(self, model, documents, request): + """Test basic two-level taxonomy creation.""" + topic_model, nr_topics = self._get_model_and_nr_topics(model, request) + if nr_topics < 4: + pytest.skip("Not enough topics for two-level taxonomy") + + has_outliers = topic_model._outliers == 1 + level_1 = max(nr_topics // 3, 2) + level_2 = max(level_1 // 3, 2) + if level_2 >= level_1: + level_2 = level_1 - 1 + + hierarchy = topic_model.create_topic_taxonomy(documents, nr_topics_per_level=[level_1, level_2]) + + assert isinstance(hierarchy, pd.DataFrame) + assert set(hierarchy.columns) == { + "Topic_ID", + "Topic_Name", + "Level", + "Parent_ID", + "Parent_Name", + "Document_Count", + } + + # Check levels exist + assert set(hierarchy.Level.unique()) == {0, 1, 2} + + # Leaf topics include non-outlier topics (and outlier if present) + leaves = hierarchy[hierarchy.Level == 0] + expected_leaf_ids = sorted([t for t in topic_model.topic_sizes_.keys() if t != -1]) + if has_outliers: + expected_leaf_ids = [-1, *expected_leaf_ids] + assert sorted(leaves.Topic_ID.tolist()) == expected_leaf_ids + + # Top-level topics have Parent_ID == -2 + top_level = hierarchy[hierarchy.Level == 2] + assert (top_level.Parent_ID == -2).all() + + # Non-top-level topics have valid parent references + for _, row in hierarchy[hierarchy.Level < 2].iterrows(): + parent_exists = (hierarchy.Topic_ID == row.Parent_ID).any() + assert parent_exists, f"Topic {row.Topic_ID} references non-existent parent {row.Parent_ID}" + + def test_single_level(self, model, documents, request): + """Test single-level taxonomy (just parents above leaves).""" + topic_model, nr_topics = self._get_model_and_nr_topics(model, request) + if nr_topics < 3: + pytest.skip("Not enough topics") + + target = max(nr_topics // 2, 2) + hierarchy = topic_model.create_topic_taxonomy(documents, nr_topics_per_level=[target]) + + assert set(hierarchy.Level.unique()) == {0, 1} + top_level = hierarchy[hierarchy.Level == 1] + assert (top_level.Parent_ID == -2).all() + + def test_document_count_consistency(self, model, documents, request): + """Test that document counts sum correctly up the hierarchy (non-outlier parents).""" + topic_model, nr_topics = self._get_model_and_nr_topics(model, request) + if nr_topics < 4: + pytest.skip("Not enough topics") + + level_1 = max(nr_topics // 2, 2) + hierarchy = topic_model.create_topic_taxonomy(documents, nr_topics_per_level=[level_1]) + + # For each non-outlier parent, doc count should equal the sum of children's doc counts + parents = hierarchy[(hierarchy.Level == 1) & (hierarchy.Topic_Name != "Outlier")] + for _, parent in parents.iterrows(): + children = hierarchy[hierarchy.Parent_ID == parent.Topic_ID] + assert parent.Document_Count == children.Document_Count.sum(), ( + f"Parent {parent.Topic_ID} doc count {parent.Document_Count} != " + f"children sum {children.Document_Count.sum()}" + ) + + def test_min_children_enforced(self, model, documents, request): + """Test that min_children constraint is enforced for non-outlier parents.""" + topic_model, nr_topics = self._get_model_and_nr_topics(model, request) + if nr_topics < 6: + pytest.skip("Not enough topics for min_children test") + + target = max(nr_topics // 2, 3) + hierarchy = topic_model.create_topic_taxonomy(documents, nr_topics_per_level=[target], min_children=2) + + # Every non-outlier parent must have at least 2 children + parents = hierarchy[(hierarchy.Level == 1) & (hierarchy.Topic_Name != "Outlier")] + for _, parent in parents.iterrows(): + n_children = (hierarchy.Parent_ID == parent.Topic_ID).sum() + assert n_children >= 2, f"Parent {parent.Topic_ID} has only {n_children} children" + + def test_min_children_three(self, model, documents, request): + """Test min_children=3 enforcement for non-outlier parents.""" + topic_model, nr_topics = self._get_model_and_nr_topics(model, request) + if nr_topics < 9: + pytest.skip("Not enough topics for min_children=3 test") + + target = max(nr_topics // 3, 3) + hierarchy = topic_model.create_topic_taxonomy(documents, nr_topics_per_level=[target], min_children=3) + + parents = hierarchy[(hierarchy.Level == 1) & (hierarchy.Topic_Name != "Outlier")] + for _, parent in parents.iterrows(): + n_children = (hierarchy.Parent_ID == parent.Topic_ID).sum() + assert n_children >= 3, f"Parent {parent.Topic_ID} has only {n_children} children (expected >= 3)" + + def test_with_embeddings(self, model, documents, document_embeddings, request): + """Test taxonomy creation with document embeddings.""" + topic_model, nr_topics = self._get_model_and_nr_topics(model, request) + if nr_topics < 3: + pytest.skip("Not enough topics") + + has_outliers = topic_model._outliers == 1 + expected_leaves = nr_topics + (1 if has_outliers else 0) + + target = max(nr_topics // 2, 2) + hierarchy = topic_model.create_topic_taxonomy( + documents, + nr_topics_per_level=[target], + embeddings=document_embeddings, + doc_embedding_weight=0.7, + ) + + assert isinstance(hierarchy, pd.DataFrame) + assert len(hierarchy[hierarchy.Level == 0]) == expected_leaves + + def test_use_ctfidf(self, model, documents, request): + """Test taxonomy creation using c-TF-IDF for distance computation.""" + topic_model, nr_topics = self._get_model_and_nr_topics(model, request) + if nr_topics < 3: + pytest.skip("Not enough topics") + + target = max(nr_topics // 2, 2) + hierarchy = topic_model.create_topic_taxonomy( + documents, + nr_topics_per_level=[target], + use_ctfidf=True, + ) + + assert isinstance(hierarchy, pd.DataFrame) + assert set(hierarchy.Level.unique()) == {0, 1} + + def test_every_leaf_has_parent(self, model, documents, request): + """Test that every leaf topic is assigned exactly one parent.""" + topic_model, nr_topics = self._get_model_and_nr_topics(model, request) + if nr_topics < 3: + pytest.skip("Not enough topics") + + target = max(nr_topics // 2, 2) + hierarchy = topic_model.create_topic_taxonomy(documents, nr_topics_per_level=[target]) + + leaves = hierarchy[hierarchy.Level == 0] + # Every leaf must have a non-null Parent_ID that is not -2 + assert leaves.Parent_ID.notna().all() + assert (leaves.Parent_ID != -2).all() + + def test_no_side_effects(self, model, documents, request): + """Test that create_topic_taxonomy does not modify the fitted model.""" + topic_model, nr_topics = self._get_model_and_nr_topics(model, request) + if nr_topics < 3: + pytest.skip("Not enough topics") + + topics_before = topic_model.topics_.copy() + sizes_before = dict(topic_model.topic_sizes_) + embeddings_before = topic_model.topic_embeddings_.copy() + + target = max(nr_topics // 2, 2) + topic_model.create_topic_taxonomy(documents, nr_topics_per_level=[target]) + + assert topic_model.topics_ == topics_before + assert dict(topic_model.topic_sizes_) == sizes_before + np.testing.assert_array_equal(topic_model.topic_embeddings_, embeddings_before) + + def test_topic_names_not_empty(self, model, documents, request): + """Test that all topics have non-empty names.""" + topic_model, nr_topics = self._get_model_and_nr_topics(model, request) + if nr_topics < 3: + pytest.skip("Not enough topics") + + target = max(nr_topics // 2, 2) + hierarchy = topic_model.create_topic_taxonomy(documents, nr_topics_per_level=[target]) + + assert (hierarchy.Topic_Name.str.len() > 0).all() + # Parent names should also be populated for non-top-level + non_top = hierarchy[hierarchy.Parent_ID != -2] + assert (non_top.Parent_Name.str.len() > 0).all() + + def test_outlier_chain(self, model, documents, request): + """Test that outlier topics form a single-child chain through all levels.""" + topic_model, nr_topics = self._get_model_and_nr_topics(model, request) + if nr_topics < 4: + pytest.skip("Not enough topics") + if topic_model._outliers != 1: + pytest.skip("No outlier topic in this model") + + level_1 = max(nr_topics // 3, 2) + level_2 = max(level_1 // 3, 2) + if level_2 >= level_1: + level_2 = level_1 - 1 + + hierarchy = topic_model.create_topic_taxonomy(documents, nr_topics_per_level=[level_1, level_2]) + + # Outlier leaf (-1) should exist + outlier_leaf = hierarchy[(hierarchy.Level == 0) & (hierarchy.Topic_ID == -1)] + assert len(outlier_leaf) == 1 + + # Walk up the chain: each outlier parent should have exactly 1 child (the outlier below) + current_id = outlier_leaf.iloc[0].Parent_ID + for level in range(1, 3): + outlier_node = hierarchy[(hierarchy.Level == level) & (hierarchy.Topic_ID == current_id)] + assert len(outlier_node) == 1 + assert outlier_node.iloc[0].Topic_Name == "Outlier" + n_children = (hierarchy.Parent_ID == current_id).sum() + assert n_children == 1 # outlier parents always have exactly 1 child + current_id = outlier_node.iloc[0].Parent_ID + + # Top-level outlier should point to Root + assert current_id == -2 + + +class TestCreateTopicTaxonomyValidation: + """Tests for parameter validation.""" + + def test_invalid_nr_topics_per_level_empty(self, base_topic_model, documents): + topic_model = copy.deepcopy(base_topic_model) + with pytest.raises(ValueError, match="non-empty list"): + topic_model.create_topic_taxonomy(documents, nr_topics_per_level=[]) + + def test_invalid_nr_topics_per_level_too_large(self, base_topic_model, documents): + topic_model = copy.deepcopy(base_topic_model) + nr_topics = len([t for t in topic_model.topic_sizes_.keys() if t != -1]) + with pytest.raises(ValueError, match="must be less than"): + topic_model.create_topic_taxonomy(documents, nr_topics_per_level=[nr_topics + 1]) + + def test_invalid_nr_topics_per_level_zero(self, base_topic_model, documents): + topic_model = copy.deepcopy(base_topic_model) + with pytest.raises(ValueError, match=">= 1"): + topic_model.create_topic_taxonomy(documents, nr_topics_per_level=[0]) + + def test_invalid_min_children(self, base_topic_model, documents): + topic_model = copy.deepcopy(base_topic_model) + with pytest.raises(ValueError, match="min_children"): + topic_model.create_topic_taxonomy(documents, nr_topics_per_level=[2], min_children=0) + + def test_invalid_doc_embedding_weight(self, base_topic_model, documents): + topic_model = copy.deepcopy(base_topic_model) + with pytest.raises(ValueError, match="doc_embedding_weight"): + topic_model.create_topic_taxonomy(documents, nr_topics_per_level=[2], doc_embedding_weight=1.5) + + def test_invalid_level_sequence(self, base_topic_model, documents): + """Test that level k+1 target must be less than level k target.""" + topic_model = copy.deepcopy(base_topic_model) + nr_topics = len([t for t in topic_model.topic_sizes_.keys() if t != -1]) + if nr_topics < 6: + pytest.skip("Not enough topics") + # level 1 = 3, level 2 = 5 is invalid (5 >= 3) + with pytest.raises(ValueError, match="must be less than"): + topic_model.create_topic_taxonomy(documents, nr_topics_per_level=[3, 5])