Chapter 5 · Text Clustering and Topic Modeling

Chapter goals: Cluster documents using multiple language models; master the complete text clustering pipeline "embeddings → dimensionality reduction → clustering → topic representation"; and upgrade it to topic modeling using BERTopic.

5.1 Preparation

💡 NOTE: Running the examples in this chapter requires a GPU. If using Google Colab, switch to Runtime > Change runtime type > Hardware accelerator > GPU > GPU type > T4.

In Colab (or other cloud environments), you need to uncomment and run the code block below to install the dependencies for this chapter:

# %%capture
# !pip install bertopic datasets openai datamapplot

5.2 ArXiv Paper Dataset: Computation and Language

We use the ArXiv paper abstract dataset from Hugging Face as the clustering corpus for this chapter:

# Load data from huggingface
from datasets import load_dataset
dataset = load_dataset("maartengr/arxiv_nlp")["train"]

# Extract metadata
abstracts = list(dataset["Abstracts"])
titles = list(dataset["Titles"])

5.3 General Pipeline for Text Clustering

5.3.1 Embedding Documents

Before clustering, convert each document into a vector representation:

from sentence_transformers import SentenceTransformer

# Create an embedding for each abstract
embedding_model = SentenceTransformer('thenlper/gte-small')
embeddings = embedding_model.encode(abstracts, show_progress_bar=True)

Inspect the dimensions of the resulting embeddings:

# Check the dimensions of the resulting embeddings
embeddings.shape

5.3.2 Reducing Embedding Dimensionality

High-dimensional embeddings are not ideal for density-based clustering; first reduce dimensions with UMAP:

from umap import UMAP

# We reduce the input embeddings from 384 dimensions to 5 dimensions
umap_model = UMAP(
    n_components=5, min_dist=0.0, metric='cosine', random_state=42
)
reduced_embeddings = umap_model.fit_transform(embeddings)

5.3.3 Clustering the Reduced Embeddings

Use HDBSCAN, a density-based algorithm that can also automatically identify outliers:

from hdbscan import HDBSCAN

# We fit the model and extract the clusters
hdbscan_model = HDBSCAN(
    min_cluster_size=50, metric='euclidean', cluster_selection_method='eom'
).fit(reduced_embeddings)
clusters = hdbscan_model.labels_

# How many clusters did we generate?
len(set(clusters))

5.4 Inspecting Clustering Results

Manually inspect the first three documents in cluster 0:

import numpy as np

# Print first three documents in cluster 0
cluster = 0
for index in np.where(clusters==cluster)[0][:3]:
    print(abstracts[index][:300] + "... \n")

Next, reduce the embeddings to 2D for plotting and a rough visual inspection of the generated clusters:

import pandas as pd

# Reduce 384-dimensional embeddings to 2 dimensions for easier visualization
reduced_embeddings = UMAP(
    n_components=2, min_dist=0.0, metric='cosine', random_state=42
).fit_transform(embeddings)

# Create dataframe
df = pd.DataFrame(reduced_embeddings, columns=["x", "y"])
df["title"] = titles
df["cluster"] = [str(c) for c in clusters]

# Select outliers and non-outliers (clusters)
clusters_df = df.loc[df.cluster != "-1", :]
outliers_df = df.loc[df.cluster == "-1", :]

Static Plot

import matplotlib.pyplot as plt

# Plot outliers and non-outliers separately
plt.scatter(outliers_df.x, outliers_df.y, alpha=0.05, s=2, c="grey")
plt.scatter(
    clusters_df.x, clusters_df.y, c=clusters_df.cluster.astype(int),
    alpha=0.6, s=2, cmap='tab20b'
)
plt.axis('off')
# plt.savefig("matplotlib.png", dpi=300)  # Uncomment to save the graph as a .png

5.5 From Text Clustering to Topic Modeling

BERTopic: A Modular Topic Modeling Framework

Assemble the three models trained above into BERTopic:

from bertopic import BERTopic

# Train our model with our previously defined models
topic_model = BERTopic(
    embedding_model=embedding_model,
    umap_model=umap_model,
    hdbscan_model=hdbscan_model,
    verbose=True
).fit(abstracts, embeddings)

Now start exploring the topics obtained above. First look at the topic overview:

topic_model.get_topic_info()

The default model generated hundreds of topics! To view the top 10 keywords for each topic along with their c-TF-IDF weights, use the get_topic() function:

topic_model.get_topic(0)

We can use the find_topics() function to search for related topics by keyword. For example, search for topics about topic modeling:

topic_model.find_topics("topic modeling")

It returns that topic 22 has high similarity (0.95) to the search term. Inspecting that topic reveals it is indeed about topic modeling:

topic_model.get_topic(22)

This topic is partially characterized by the classic LDA technique. Let's check whether the BERTopic paper itself was also assigned to topic 22:

topic_model.topics_[titles.index('BERTopic: Neural topic modeling with a class-based TF-IDF procedure')]

Yes! This matches expectations, since words like "clustering" and "topic" that are not LDA-specific appear in the topic description.

Visualization

Document Visualization

# Visualize topics and documents
fig = topic_model.visualize_documents(
    titles,
    reduced_embeddings=reduced_embeddings,
    width=1200,
    hide_annotations=True
)

# Update fonts of legend for easier visualization
fig.update_layout(font=dict(size=16))

# Visualize barchart with ranked keywords
topic_model.visualize_barchart()

# Visualize relationships between topics
topic_model.visualize_heatmap(n_clusters=30)

# Visualize the potential hierarchical structure of topics
topic_model.visualize_hierarchy()

5.6 Representation Models

In the example below, we update topic representations after the model has finished training, allowing for rapid iteration. However, if you want to use a representation model from the start of training, run it like this:

from bertopic.representation import KeyBERTInspired
from bertopic import BERTopic

# Create your representation model
representation_model = KeyBERTInspired()

# Use the representation model in BERTopic on top of the default pipeline
topic_model = BERTopic(representation_model=representation_model)

For easy comparison between with-and-without representation models, first copy the current topic model's original representations:

# Save original representations
from copy import deepcopy
original_topics = deepcopy(topic_model.topic_representations_)

Define a utility function for displaying the difference in topic representations between two models:

def topic_differences(model, original_topics, nr_topics=5):
    """Show the differences in topic representations between two models """
    df = pd.DataFrame(columns=["Topic", "Original", "Updated"])
    for topic in range(nr_topics):

        # Extract top 5 words per topic per model
        og_words = " | ".join(list(zip(*original_topics[topic]))[0][:5])
        new_words = " | ".join(list(zip(*model.get_topic(topic)))[0][:5])
        df.loc[len(df)] = [topic, og_words, new_words]

    return df

KeyBERTInspired

from bertopic.representation import KeyBERTInspired

# Update our topic representations using KeyBERTInspired
representation_model = KeyBERTInspired()
topic_model.update_topics(abstracts, representation_model=representation_model)

# Show topic differences
topic_differences(topic_model, original_topics)

Maximal Marginal Relevance

from bertopic.representation import MaximalMarginalRelevance

# Update our topic representations to MaximalMarginalRelevance
representation_model = MaximalMarginalRelevance(diversity=0.5)
topic_model.update_topics(abstracts, representation_model=representation_model)

# Show topic differences
topic_differences(topic_model, original_topics)

5.7 Text Generation

Flan-T5

from transformers import pipeline
from bertopic.representation import TextGeneration

prompt = """I have a topic that contains the following documents:
[DOCUMENTS]

The topic is described by the following keywords: '[KEYWORDS]'.

Based on the documents and keywords, what is this topic about?"""

# Update our topic representations using Flan-T5
generator = pipeline('text2text-generation', model='google/flan-t5-small')
representation_model = TextGeneration(
    generator, prompt=prompt, doc_length=50, tokenizer="whitespace"
)
topic_model.update_topics(abstracts, representation_model=representation_model)

# Show topic differences
topic_differences(topic_model, original_topics)

OpenAI

import openai
from bertopic.representation import OpenAI

prompt = """
I have a topic that contains the following documents:
[DOCUMENTS]

The topic is described by the following keywords: [KEYWORDS]

Based on the information above, extract a short topic label in the following format:
topic: <short topic label>
"""

# Update our topic representations using GPT-3.5
client = openai.OpenAI(api_key="YOUR_KEY_HERE")
representation_model = OpenAI(
    client, model="gpt-3.5-turbo", exponential_backoff=True, chat=True, prompt=prompt
)
topic_model.update_topics(abstracts, representation_model=representation_model)

# Show topic differences
topic_differences(topic_model, original_topics)

Use DatamapPlot for the final visualization of topics and documents:

# Visualize topics and documents
fig = topic_model.visualize_document_datamap(
    titles,
    topics=list(range(20)),
    reduced_embeddings=reduced_embeddings,
    width=1200,
    label_font_size=11,
    label_wrap_width=20,
    use_medoids=True,
)
plt.savefig("datamapplot.png", dpi=300)

5.8 BONUS: Word Clouds

Before following this section, please pip install wordcloud first.

First, let each topic be described by more than 10 words — this makes the generated word clouds much more interesting:

topic_model.update_topics(abstracts, top_n_words=500)

Then run the following code to generate word clouds for our topics:

from wordcloud import WordCloud
import matplotlib.pyplot as plt

def create_wordcloud(model, topic):
    plt.figure(figsize=(10,5))
    text = {word: value for word, value in model.get_topic(topic)}
    wc = WordCloud(background_color="white", max_words=1000, width=1600, height=800)
    wc.generate_from_frequencies(text)
    plt.imshow(wc, interpolation="bilinear")
    plt.axis("off")
    plt.show()

# Show wordcloud
create_wordcloud(topic_model, topic=17)

Chapter Summary

• The general text clustering pipeline has three core steps: SentenceTransformer embeds documents → UMAP reduces dimensionality → HDBSCAN density clustering;
• HDBSCAN's min_cluster_size controls minimum cluster size; the -1 label represents outliers;
• BERTopic is a modular framework that can reuse custom-built embedding, UMAP, and HDBSCAN models; topic keywords are weighted by c-TF-IDF;
• Representation models (KeyBERTInspired, MaximalMarginalRelevance, TextGeneration, etc.) can be rapidly iterated after training via update_topics() to refine topic representations;
• Built-in visualization methods such as visualize_documents(), visualize_barchart(), and visualize_heatmap() help explore topic structure and relationships.

🧪 随堂测验

点击你认为正确的选项。答错时会展示正确答案与原因解析。

1. What are the three core steps of the text clustering pipeline in this chapter, in order?

2. What do documents with label -1 in the HDBSCAN output represent?

3. What is the role of c-TF-IDF in BERTopic?

4. To rapidly try different topic representation schemes after the model has finished training, which method should you use?

🛠️ Hands-on Practice

1. Change HDBSCAN's min_cluster_size to 20 and 100 respectively; observe how the number of generated clusters changes; and use the scatter-plot code from Section 5.4 to compare clustering shapes across the three settings.
2. Use find_topics() to search for a research direction you are interested in (e.g. "reinforcement learning" or "image segmentation"), find the corresponding topic number, print the top 10 keywords for that topic, and manually judge relevance.
3. Following the OpenAI example in Section 5.7, swap the model parameter with another Chat model you have access to (e.g. DeepSeek); regenerate labels for the top 5 topics and compare the difference table against original_topics.