BERTopic建模时如何获取单个主题对应的全部文档

问题原因 #

get_representative_docs()接口的设计目标是返回每个主题最具代表性的核心样本，默认固定返回3条，本身就不支持拉取对应主题下的全量文档。

解决方法 #

你调用fit_transform()时返回的topics数组和输入的docs_test语料是严格一一对应的，数组中每个位置的数值就是对应位置文档所属的主题编号，基于这个映射关系就可以筛选出任意主题下的全部文档。

方法1：直接遍历筛选

# 定义你要查询的目标主题编号
target_topic = 1
# 筛选属于目标主题的所有文档
all_docs_of_topic = [doc for doc, topic_num in zip(docs_test, topics) if topic_num == target_topic]
# 验证文档数量
print(f"主题{target_topic}总文档数：{len(all_docs_of_topic)}")

方法2：转换为DataFrame方便后续批量查询

如果需要频繁查询不同主题的文档，可以先把映射关系存为DataFrame：

import pandas as pd

# 生成文档-主题对应表
doc_topic_mapping = pd.DataFrame({
    "document_content": docs_test,
    "topic_id": topics
})

# 筛选指定主题的全量文档
target_topic = 1
all_docs_of_topic = doc_topic_mapping[doc_topic_mapping["topic_id"] == target_topic]["document_content"].tolist()

注意事项 #

如果你的模型配置了自动合并主题，或者存在标记为-1的离群文档（未被划分到任何有效主题的样本），也可以用同样的筛选逻辑获取对应文档。

内容的提问来源于stack exchange，提问作者Kaleem