Skip to content

第 5 章 · 文本聚类与主题建模

本章目标:使用多种语言模型对文档进行聚类,掌握「嵌入 → 降维 → 聚类 → 主题表示」的完整文本聚类流水线,并用 BERTopic 将其升级为主题建模。

5.1 准备工作

💡 提示:运行本章示例需要 GPU。如果使用 Google Colab,请在 Runtime > Change runtime type > Hardware accelerator > GPU > GPU type > T4 中切换。

在 Colab(或其他云环境)中需要取消注释并运行以下代码块来安装本章依赖:

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

5.2 ArXiv 文章数据集:计算与语言

我们使用 Hugging Face 上的 ArXiv 论文摘要数据集作为本章的聚类语料:

python
# 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 文本聚类的通用流程

5.3.1 嵌入文档

聚类之前,先把每篇文档转换成向量表示:

python
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)

检查生成的嵌入维度:

python
# Check the dimensions of the resulting embeddings
embeddings.shape

5.3.2 降低嵌入的维度

高维嵌入不利于密度聚类,先用 UMAP 降维:

python
from umap import UMAP

# We reduce the input embeddings from 384 dimenions to 5 dimenions
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 对降维后的嵌入进行聚类

使用 HDBSCAN 这种基于密度的算法,它还能自动识别离群点:

python
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 检查聚类结果

手动查看第 0 个簇中的前三篇文档:

python
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")

接下来把嵌入降到二维,以便绘图并粗略观察生成的簇:

python
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", :]

静态图

python
import matplotlib.pyplot as plt

# Plot outliers and non-outliers seperately
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 从文本聚类到主题建模

BERTopic:模块化主题建模框架

把前面训练好的三个模型组装进 BERTopic:

python
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)

现在开始探索上面得到的主题。先看主题总览:

python
topic_model.get_topic_info()

默认模型生成了数百个主题!要查看每个主题的前 10 个关键词及其 c-TF-IDF 权重,可以使用 get_topic() 函数:

python
topic_model.get_topic(0)

我们可以用 find_topics() 函数根据搜索词查找相关主题。比如搜索关于主题建模(topic modeling)的主题:

python
topic_model.find_topics("topic modeling")

它返回主题 22 与搜索词有较高相似度(0.95)。查看该主题,可以发现它确实是关于主题建模的:

python
topic_model.get_topic(22)

这个主题部分上由经典的 LDA 技术刻画。看看 BERTopic 论文本身是否也被分配到了主题 22:

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

是的!这符合预期,因为主题描述中出现了 "clustering"、"topic" 等非 LDA 特有的词。

可视化

文档可视化

python
# 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))
python
# 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 表示模型

在下面的例子中,我们会在模型训练完成之后更新主题表示,这样可以快速迭代。不过,如果你想在训练一开始就使用表示模型,可以这样运行:

python
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)

为了方便对比有无表示模型的差异,先复制一份当前主题模型的原有表示:

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

定义一个用于展示两种模型主题表示差异的工具函数:

python
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

python
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)

python
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 文本生成

Flan-T5

python
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

python
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)

用 DatamapPlot 对主题和文档做最终的可视化展示:

python
# 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:词云

跟随本节操作前,请先 pip 安装 wordcloud

首先,让每个主题由比 10 个更多的词来描述——这样生成的词云会有趣得多:

python
topic_model.update_topics(abstracts, top_n_words=500)

然后运行以下代码为我们的主题生成词云:

python
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)

本章小结

  • 文本聚类通用流水线三步走:SentenceTransformer 嵌入文档 → UMAP 降维 → HDBSCAN 密度聚类;
  • HDBSCAN 的 min_cluster_size 控制最小簇大小,-1 标签代表离群点;
  • BERTopic 是模块化框架,可以复用自建的 embedding、UMAP、HDBSCAN 模型,主题关键词由 c-TF-IDF 加权;
  • 表示模型(KeyBERTInspired、MaximalMarginalRelevance、TextGeneration 等)可在训练后通过 update_topics() 快速迭代优化主题表示;
  • visualize_documents()visualize_barchart()visualize_heatmap() 等内置可视化方法帮助探索主题结构与关系。

🧪 随堂测验

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

1. 本章文本聚类流水线的三个核心步骤按顺序是?

2. HDBSCAN 输出中标签为 -1 的文档代表什么?

3. BERTopic 中 c-TF-IDF 的作用是?

4. 想在模型训练完成之后快速尝试不同的主题表示方案,应该使用哪个方法?

🛠️ 动手实践

  1. 把 HDBSCAN 的 min_cluster_size 分别改为 20 和 100,观察生成的簇数量变化,并用 5.4 节的散点图代码对比三种设置下的聚类形态。
  2. find_topics() 搜索一个你感兴趣的研究方向(如 "reinforcement learning" 或 "image segmentation"),找到对应主题编号,打印该主题的前 10 个关键词并人工判断相关性。
  3. 参照 5.7 节的 OpenAI 示例,把 model 参数换成你可用的其他 Chat 模型(如 DeepSeek),重新生成前 5 个主题的标签,并与 original_topics 的差异表对比效果。