Skip to content

Commit cbbb473

Browse files
committed
Add create_topic_taxonomy() method to _bertopic (#2526)
1 parent f969760 commit cbbb473

4 files changed

Lines changed: 741 additions & 0 deletions

File tree

‎bertopic/_bertopic.py‎

Lines changed: 312 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1201,6 +1201,318 @@ def hierarchical_topics(
12011201

12021202
return hier_topics
12031203

1204+
def create_topic_taxonomy(
1205+
self,
1206+
docs: List[str],
1207+
nr_topics_per_level: List[int],
1208+
embeddings: np.ndarray = None,
1209+
min_children: int = 2,
1210+
doc_embedding_weight: float = 0.5,
1211+
use_ctfidf: bool = False,
1212+
) -> pd.DataFrame:
1213+
"""Builds a deterministic N-level taxonomy on top of the fitted leaf topics.
1214+
Each level is created by clustering the level below using agglomerative
1215+
clustering with cosine distance, then repairing any clusters that have
1216+
fewer than ``min_children`` by merging them into the most similar neighbor.
1217+
1218+
The leaf-level topic model remains unchanged - parent topics are derived
1219+
through a lookup table and are not used during inference.
1220+
1221+
Arguments:
1222+
docs: The documents you used when calling either ``fit`` or ``fit_transform``
1223+
nr_topics_per_level: Target number of topics for each level above the leaves.
1224+
For example, ``[20, 5]`` creates two parent levels: level 1 with ~20 topics
1225+
and level 2 with ~5 topics. Each entry must be less than the number of
1226+
topics at the level below.
1227+
embeddings: Pre-trained document embeddings used to compute document centroids
1228+
per topic. If ``None``, only BERTopic's internal topic embeddings are used
1229+
and ``doc_embedding_weight`` is ignored.
1230+
min_children: Minimum number of children each parent must have at every level.
1231+
Clusters with fewer children are merged into their most similar neighbor.
1232+
Must be >= 1. Default is 2.
1233+
doc_embedding_weight: Weight for document centroids when computing leaf topic
1234+
vectors. The topic embedding weight is ``1 - doc_embedding_weight``.
1235+
Only used when ``embeddings`` is provided. Must be in [0, 1]. Default is 0.5.
1236+
use_ctfidf: Whether to use c-TF-IDF representations for clustering distance
1237+
computation. If ``False``, topic embedding vectors are used. Default is ``False``.
1238+
1239+
Returns:
1240+
hierarchy: A DataFrame with the full hierarchy. Columns:
1241+
1242+
- ``Topic_ID``: Unique topic identifier. Leaf IDs match the fitted model's
1243+
topic IDs. Parent IDs are auto-incremented integers.
1244+
- ``Topic_Name``: Top-5 c-TF-IDF words joined by ``_``.
1245+
- ``Level``: Hierarchy level (0 = leaf, 1 = first parent, etc.).
1246+
- ``Parent_ID``: ID of the parent topic at the next level up. ``-2`` for
1247+
top-level topics.
1248+
- ``Parent_Name``: Name of the parent topic. ``"Root"`` for top-level.
1249+
- ``Document_Count``: Number of documents. For parents this is the sum of
1250+
their children's counts.
1251+
1252+
Examples:
1253+
```python
1254+
from bertopic import BERTopic
1255+
topic_model = BERTopic()
1256+
topics, probs = topic_model.fit_transform(docs)
1257+
1258+
# Two-level hierarchy: ~20 parents, ~5 grandparents
1259+
hierarchy = topic_model.create_topic_taxonomy(
1260+
docs, nr_topics_per_level=[20, 5], min_children=2
1261+
)
1262+
1263+
# With document embeddings for richer topic vectors
1264+
hierarchy = topic_model.create_topic_taxonomy(
1265+
docs, nr_topics_per_level=[20, 5],
1266+
embeddings=embeddings, doc_embedding_weight=0.7
1267+
)
1268+
```
1269+
"""
1270+
check_is_fitted(self)
1271+
check_documents_type(docs)
1272+
1273+
# Validate parameters
1274+
if not isinstance(nr_topics_per_level, list) or len(nr_topics_per_level) == 0:
1275+
raise ValueError("`nr_topics_per_level` must be a non-empty list of positive integers.")
1276+
if any(n < 1 for n in nr_topics_per_level):
1277+
raise ValueError("All entries in `nr_topics_per_level` must be >= 1.")
1278+
if min_children < 1:
1279+
raise ValueError("`min_children` must be >= 1.")
1280+
if not 0.0 <= doc_embedding_weight <= 1.0:
1281+
raise ValueError("`doc_embedding_weight` must be between 0.0 and 1.0.")
1282+
1283+
# Non-outlier leaf topics
1284+
topic_ids = sorted([t for t in self.topic_sizes_.keys() if t != -1])
1285+
n_leaves = len(topic_ids)
1286+
if n_leaves < 2:
1287+
raise ValueError("Cannot create hierarchy with fewer than 2 non-outlier topics.")
1288+
1289+
# Validate level targets against available topics
1290+
n_at_level = n_leaves
1291+
for k, target_n in enumerate(nr_topics_per_level):
1292+
if target_n >= n_at_level:
1293+
raise ValueError(
1294+
f"nr_topics_per_level[{k}] = {target_n} must be less than the "
1295+
f"number of topics at the level below ({n_at_level})."
1296+
)
1297+
n_at_level = target_n
1298+
1299+
if embeddings is not None:
1300+
check_embeddings_shape(embeddings, docs)
1301+
1302+
# Phase 1: Prepare leaf level
1303+
topic_to_idx = {tid: i for i, tid in enumerate(topic_ids)}
1304+
topic_embeds = self.topic_embeddings_[self._outliers :]
1305+
1306+
# Compute leaf topic vectors
1307+
if embeddings is not None:
1308+
doc_topics = np.array(self.topics_)
1309+
doc_centroids = np.zeros_like(topic_embeds)
1310+
for tid in topic_ids:
1311+
mask = doc_topics == tid
1312+
doc_centroids[topic_to_idx[tid]] = embeddings[mask].mean(axis=0)
1313+
w = doc_embedding_weight
1314+
level_vectors = w * doc_centroids + (1.0 - w) * topic_embeds
1315+
else:
1316+
level_vectors = topic_embeds.copy()
1317+
1318+
doc_counts = np.array([self.topic_sizes_[tid] for tid in topic_ids])
1319+
1320+
# Prepare BoW for label generation
1321+
documents = pd.DataFrame({"Document": docs, "ID": range(len(docs)), "Topic": self.topics_})
1322+
documents_per_topic = documents.groupby(["Topic"], as_index=False).agg({"Document": " ".join})
1323+
documents_per_topic = documents_per_topic.loc[documents_per_topic.Topic != -1, :]
1324+
clean_documents = self._preprocess_text(documents_per_topic.Document.values)
1325+
1326+
# Support older sklearn versions, as was done in other parts as well.
1327+
if version.parse(sklearn_version) >= version.parse("1.0.0"):
1328+
words = self.vectorizer_model.get_feature_names_out()
1329+
else:
1330+
words = self.vectorizer_model.get_feature_names()
1331+
1332+
bow = self.vectorizer_model.transform(clean_documents)
1333+
1334+
# Reorder BoW rows to match topic_ids order
1335+
bow_topic_order = documents_per_topic.Topic.tolist()
1336+
bow_idx_map = {tid: i for i, tid in enumerate(bow_topic_order)}
1337+
reorder = [bow_idx_map[tid] for tid in topic_ids]
1338+
bow = bow[reorder]
1339+
1340+
# Initialize hierarchy records with leaf topics
1341+
records = []
1342+
for tid in topic_ids:
1343+
label = self.topic_labels_.get(tid, f"Topic_{tid}")
1344+
records.append(
1345+
{
1346+
"Topic_ID": tid,
1347+
"Topic_Name": label,
1348+
"Level": 0,
1349+
"Parent_ID": None,
1350+
"Parent_Name": None,
1351+
"Document_Count": self.topic_sizes_[tid],
1352+
}
1353+
)
1354+
1355+
# Track current level state
1356+
current_ids = list(topic_ids)
1357+
current_vectors = level_vectors
1358+
current_doc_counts = doc_counts
1359+
current_bow = bow
1360+
next_parent_id = max(topic_ids) + 1
1361+
# Track which original leaf topic IDs each node covers (for BoW/doc lookups)
1362+
current_leaf_sets = [[tid] for tid in topic_ids]
1363+
1364+
# Phase 2: Build each parent level
1365+
for k, target_n in enumerate(nr_topics_per_level, start=1):
1366+
n_current = len(current_ids)
1367+
1368+
# Clamp target to actual available topics (prior-level repair may have reduced the count)
1369+
if target_n >= n_current:
1370+
target_n = n_current - 1
1371+
if target_n < 1:
1372+
break
1373+
1374+
# Compute distance matrix
1375+
if use_ctfidf:
1376+
c_tf_idf_level = self.ctfidf_model.transform(current_bow)
1377+
if isinstance(c_tf_idf_level, csr_matrix):
1378+
c_tf_idf_level = c_tf_idf_level.toarray()
1379+
dist_matrix = 1 - cosine_similarity(c_tf_idf_level)
1380+
else:
1381+
dist_matrix = 1 - cosine_similarity(current_vectors)
1382+
np.fill_diagonal(dist_matrix, 0)
1383+
1384+
# Agglomerative clustering (again, support older versions)
1385+
if version.parse(sklearn_version) >= version.parse("1.4.0"):
1386+
agg = AgglomerativeClustering(target_n, metric="precomputed", linkage="average")
1387+
else:
1388+
agg = AgglomerativeClustering(target_n, affinity="precomputed", linkage="average")
1389+
agg.fit(dist_matrix)
1390+
1391+
# Build cluster-to-children mapping
1392+
cluster_to_children = defaultdict(list)
1393+
for i, cl in enumerate(agg.labels_):
1394+
cluster_to_children[cl].append(i)
1395+
1396+
# "Repair": merge clusters with fewer than min_children
1397+
changed = True
1398+
while changed:
1399+
changed = False
1400+
centroids = {}
1401+
for cl, children_idxs in cluster_to_children.items():
1402+
weights = current_doc_counts[children_idxs]
1403+
centroids[cl] = np.average(current_vectors[children_idxs], axis=0, weights=weights)
1404+
1405+
small_clusters = [cl for cl, ch in cluster_to_children.items() if len(ch) < min_children]
1406+
for cl in small_clusters:
1407+
if cl not in cluster_to_children or len(cluster_to_children) <= 1:
1408+
continue
1409+
1410+
# Find most similar neighbor by cosine similarity
1411+
cl_centroid = centroids[cl].reshape(1, -1)
1412+
best_neighbor = None
1413+
best_sim = -1.0
1414+
for other_cl, other_centroid in centroids.items():
1415+
if other_cl == cl:
1416+
continue
1417+
sim = cosine_similarity(cl_centroid, other_centroid.reshape(1, -1))[0, 0]
1418+
if sim > best_sim:
1419+
best_sim = sim
1420+
best_neighbor = other_cl
1421+
1422+
# Merge into best neighbor
1423+
cluster_to_children[best_neighbor].extend(cluster_to_children[cl])
1424+
del cluster_to_children[cl]
1425+
# Recompute merged centroid
1426+
idxs = cluster_to_children[best_neighbor]
1427+
weights = current_doc_counts[idxs]
1428+
centroids[best_neighbor] = np.average(current_vectors[idxs], axis=0, weights=weights)
1429+
del centroids[cl]
1430+
changed = True
1431+
1432+
# Build parent nodes for this level
1433+
next_level_ids = []
1434+
next_level_vectors = []
1435+
next_level_doc_counts = []
1436+
next_level_bow = []
1437+
next_level_leaf_sets = []
1438+
1439+
for cl in sorted(cluster_to_children.keys()):
1440+
children_idxs = cluster_to_children[cl]
1441+
parent_id = next_parent_id
1442+
next_parent_id += 1
1443+
1444+
# Collect all original leaf topic IDs covered by this parent
1445+
parent_leaf_ids = []
1446+
for i in children_idxs:
1447+
parent_leaf_ids.extend(current_leaf_sets[i])
1448+
1449+
# Compute parent label via c-TF-IDF on merged BoW
1450+
parent_bow = csr_matrix(current_bow[children_idxs].sum(axis=0))
1451+
c_tf_idf_parent = self.ctfidf_model.transform(parent_bow)
1452+
temp_docs = documents.loc[documents.Topic.isin(parent_leaf_ids)].copy()
1453+
temp_docs["Topic"] = 0
1454+
words_per_topic = self._extract_words_per_topic(
1455+
words,
1456+
temp_docs,
1457+
c_tf_idf_parent,
1458+
calculate_aspects=False,
1459+
fine_tune_representation=False,
1460+
)
1461+
parent_name = "_".join([x[0] for x in words_per_topic[0]][:5])
1462+
1463+
# Compute parent vector and doc count
1464+
weights = current_doc_counts[children_idxs]
1465+
parent_vector = np.average(current_vectors[children_idxs], axis=0, weights=weights)
1466+
parent_doc_count = int(current_doc_counts[children_idxs].sum())
1467+
1468+
# Update children's parent references
1469+
for idx in children_idxs:
1470+
child_id = current_ids[idx]
1471+
for rec in records:
1472+
if rec["Topic_ID"] == child_id and rec["Level"] == k - 1:
1473+
rec["Parent_ID"] = parent_id
1474+
rec["Parent_Name"] = parent_name
1475+
break
1476+
1477+
records.append(
1478+
{
1479+
"Topic_ID": parent_id,
1480+
"Topic_Name": parent_name,
1481+
"Level": k,
1482+
"Parent_ID": None,
1483+
"Parent_Name": None,
1484+
"Document_Count": parent_doc_count,
1485+
}
1486+
)
1487+
1488+
next_level_ids.append(parent_id)
1489+
next_level_vectors.append(parent_vector)
1490+
next_level_doc_counts.append(parent_doc_count)
1491+
next_level_bow.append(parent_bow)
1492+
next_level_leaf_sets.append(parent_leaf_ids)
1493+
1494+
# Prepare for next level
1495+
current_ids = next_level_ids
1496+
current_vectors = np.array(next_level_vectors)
1497+
current_doc_counts = np.array(next_level_doc_counts)
1498+
current_bow = sp.vstack(next_level_bow)
1499+
current_leaf_sets = next_level_leaf_sets
1500+
1501+
# Stop if we can't cluster further
1502+
if len(current_ids) <= 1:
1503+
break
1504+
1505+
# Phase 3: Finalize top-level parents
1506+
max_level = max(rec["Level"] for rec in records)
1507+
for rec in records:
1508+
if rec["Level"] == max_level and rec["Parent_ID"] is None:
1509+
rec["Parent_ID"] = -2
1510+
rec["Parent_Name"] = "Root"
1511+
1512+
hierarchy_df = pd.DataFrame(records)
1513+
hierarchy_df = hierarchy_df.sort_values(["Level", "Topic_ID"]).reset_index(drop=True)
1514+
return hierarchy_df
1515+
12041516
def approximate_distribution(
12051517
self,
12061518
documents: Union[str, List[str]],

0 commit comments

Comments
 (0)