如何在bertopic建模中按主题获取所有文档

Question

我有一个数据集，正在尝试使用berTopic建模将其转换为主题，但问题是，我无法获取主题的全部文档。 berTopic每个主题只返回3个文档。

topic_model  = BERTopic(verbose=True, embedding_model=embedding_model,
                                nr_topics = 'auto',
                                n_gram_range = (3,3),
                                top_n_words = 10,
                               calculate_probabilities=True, 
                              seed_topic_list = topic_list,
                              )
topics, probs = topic_model.fit_transform(docs_test)
representative_doc = topic_model.get_representative_docs(topic#1)
representative_doc

这个主题包含超过300篇文档，但使用.get_representative_docs只有3篇被bertopic显示出来。

- Kaleem

Answer 1

可能有更优雅的解决方案，但我不是专家，可以分享我所用的方法（因为还没有答案）：

"topics, probs = topic_model.fit_transform(docs_test)"返回主题。

因此，您可以将此输出与文档组合在一起。例如，使用以下方法将它们组合成(pandas.)数据帧：

df = pd.DataFrame({'topic': topics, 'document': docs_test})

现在您可以根据每个主题过滤此数据框，以识别引用文档。

topic_0 = df[df.topic == 0]

- LM-me

Answer 2

BERTopic提供了一个API get_document_info()，它可以返回每个文档及其相关主题的数据框。https://maartengr.github.io/BERTopic/api/bertopic.html#bertopic._bertopic.BERTopic.get_document_info

此API的响应如下所示:

索引	文档	主题	名称	...
0	doc1_text	241	kw1_kw2_	...
1	doc2_text	-1	kw1_kw2_	...

您可以使用此数据框通过Pandas groupby或其他方法获取与特定主题相关的所有文档。

T = topic_model.get_document_info(docs)
docs_per_topics = T.groupby(["Topic"]).apply(lambda x: x.index).to_dict()

代码返回如下所示的字典：

{
    -1: Int64Index([3,10,11,12,15,16,18,19,20,22,...365000], dtype='int64',length=149232),
    0: Int64Index([907,1281,1335,1337,...308420,308560,308645],dtype='int64',length=5127),
    ...
}

- Jay Patel

---

当Bertopic没有任何API时，它的早期版本存在这个问题，现在则非常容易解决。 - Kaleem