python - 如何在bertopic建模中获得主题概率矩阵

Question

我运行 BERTopic 来获取 3,500 个文档的主题。如何获取每个文档的主题概率矩阵并将它们导出到 csv？当我导出它们时，我也想导出每个文档的标识符。

我尝试了两种方法:首先，我发现 topic_model.visualize_distribution(probs[#]) 提供了我想要的信息。但是如何将每个文档的主题-概率数据导出到 csv？

其次，我发现如果我可以将概率列添加到它生成的数据框中，这个线程 ( How to get all docoments per topic in bertopic modeling) 会很有用。有什么办法吗？

请分享可以为所有文档生成和导出主题概率矩阵的任何其他方法。

供您引用，这是我的 BERTopic 代码。谢谢!

embedding_model = SentenceTransformer('all-mpnet-base-v2')
umap_model = UMAP(n_neighbors=15)
hdbscan_model = HDBSCAN(min_cluster_size=20, min_samples=1,
                        gen_min_span_tree=True,
                        prediction_data=True)

stopwords = list(stopwords.words('english')) + ['http', 'https', 'amp', 'com']
vectorizer_model = CountVectorizer(ngram_range=(1, 3), stop_words=stopwords)

model1 = BERTopic(
    umap_model=umap_model,
    hdbscan_model=hdbscan_model,
    embedding_model=embedding_model,
    vectorizer_model=vectorizer_model,
    language='english',
    calculate_probabilities=True,
    verbose=True
)
topics, probs = model1.fit_transform(data)

Answer 1

probs 变量包含对应于每个单独文档的所有主题概率。您可以像这样从这些值创建数据框:

#convert 2D array to pandas DataFrame    
topic_prob_df = pd.DataFrame(probs)
#create 'data' column - or, alternatively, an identifier column for data
topic_prob_df['data'] = data
#export as csv 
topic_prob_df.to_csv('topic-probs.csv')