如何在BERTopic中获取并导出带文档标识的主题概率矩阵

方法1：直接基于probs数组构建DataFrame #

你已经通过fit_transform得到了topics（每份文档的分配主题）和probs（每份文档对应所有主题的概率数组），结合文档标识（假设你有单独的doc_ids列表，或是data本身包含文档ID列），可直接构建完整的DataFrame：

import pandas as pd

# 假设你有文档标识列表doc_ids，若data是DataFrame，可通过doc_ids = data["doc_id"].tolist()获取
doc_info = pd.DataFrame({
    "doc_id": doc_ids,
    "assigned_topic": topics,
    "text": data  # 可选：保留原文方便后续核对
})

# 将probs数组转为DataFrame，列名用主题编号命名
probs_df = pd.DataFrame(probs, columns=[f"topic_{topic}_prob" for topic in model1.get_topics().keys()])

# 合并文档信息与概率矩阵
final_df = pd.concat([doc_info, probs_df], axis=1)

# 导出到CSV
final_df.to_csv("document_topic_probs.csv", index=False, encoding="utf-8")

方法2：利用BERTopic自带的get_document_info方法 #

BERTopic的get_document_info方法会返回包含文档主题、相似度等基础信息的DataFrame，只需追加概率列即可：

import pandas as pd

# 获取基础文档信息
doc_info = model1.get_document_info(data)

# 转换probs为DataFrame并添加列名
probs_df = pd.DataFrame(probs, columns=[f"topic_{topic}_prob" for topic in model1.get_topics().keys()])

# 合并后导出
final_df = pd.concat([doc_info, probs_df], axis=1)
final_df.to_csv("document_topic_probs.csv", index=False, encoding="utf-8")

针对你之前的疑问解答 #

1. visualize_distribution的数据导出：该可视化函数的数据源就是你已获取的probs数组，无需从可视化界面导出，直接用probs构建概率矩阵即可。
2. 给Stack Overflow的DataFrame添加概率列：核心操作就是用pd.concat将文档信息DataFrame和probs转换后的DataFrame按列合并，只要两者行顺序与原始文档顺序一致（probs默认与输入data顺序完全对应）即可。

额外注意事项 #

• 若data是纯文本列表，务必提前准备好对应的文档标识（如文件名、数据库ID等），否则导出的CSV无法关联到原始文档。
• 主题编号可通过model1.get_topics().keys()获取，确保概率列的主题编号与BERTopic生成的主题完全匹配。

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