我有一个数据集并尝试使用 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
可能有更优雅的解决方案,因为我不是专家,但我可以分享对我有用的方法:
topics, probs = topic_model.fit_transform(docs_test)
返回主题。
因此,您可以将此输出与文档结合起来。 例如,使用以下方法将它们组合成 Pandas 数据框:
df = pd.DataFrame({'topic': topics, 'document': docs_test})
现在,您可以针对每个主题过滤此数据框以识别引用文档:
topic_0 = df[df.topic == 0]
BERTopic
get_document_info()
有一个 API,它返回每个文档及其关联主题的数据帧。 https://maartengr.github.io/BERTopic/api/bertopic.html#bertopic._bertopic.BERTopic.get_document_info
该API的响应如下所示:
索引 | 文件 | 主题 | 姓名 | ... |
---|---|---|---|---|
0 | doc1_文本 | 241 | kw1_kw2_ | ... |
1 | doc2_文本 | -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),
...
}