gensim计算文档相似度

gensim是一个基于Python语言的开源工具集,用于处理向量空间模型(vector space modeling)和话题模型(topic modeling)的相关问题。本文分享如何使用gensim工具来计算两篇中文文档的相似度。

首先我们要生成一些中文文档。下面的代码生成一个名为documents的文档列表,由于是演示,这里的每个文档中只有几个词语。关于如何对原始的字符串做分词操作,可以参考jieba

1
2
3
4
5
6
7
8
9
10
11
from gensim import corpora
from pprint import pprint
documents = ["你好 好的 不错 笨蛋",
"笨蛋 傻瓜 傻子 哈哈",
"好的 不错 你好 哈哈",
"有趣 可以 好的 不错 还行",
"傻瓜 傻子 二货 还行",
"可以 好的 不错 哈哈",
"有趣 有趣 哈哈 哈哈"]
texts = [[word for word in document.split()] for document in documents]
pprint(texts)
[['你好', '好的', '不错', '笨蛋'],
 ['笨蛋', '傻瓜', '傻子', '哈哈'],
 ['好的', '不错', '你好', '哈哈'],
 ['有趣', '可以', '好的', '不错', '还行'],
 ['傻瓜', '傻子', '二货', '还行'],
 ['可以', '好的', '不错', '哈哈'],
 ['有趣', '有趣', '哈哈', '哈哈']]

然后我们要对原始的文档做一些预处理。这里我们统计了每个词语出现的次数,移除了只出现一次的词语,因为这些词语通常不会对计算文档相似度产生任何贡献。其他的预处理操作还包括移除数字、字母以及标点符号等,这里我们没有展示,如果实际中遇到可以进行相应的处理。

1
2
3
4
5
6
7
from collections import defaultdict
frequency = defaultdict(int)
for text in texts:
for token in text:
frequency[token] += 1
texts = [[token for token in text if frequency[token] > 1] for text in texts]
pprint(texts)
[['你好', '好的', '不错', '笨蛋'],
 ['笨蛋', '傻瓜', '傻子', '哈哈'],
 ['好的', '不错', '你好', '哈哈'],
 ['有趣', '可以', '好的', '不错', '还行'],
 ['傻瓜', '傻子', '还行'],
 ['可以', '好的', '不错', '哈哈'],
 ['有趣', '有趣', '哈哈', '哈哈']]

接着我们生成了包含所有词语的词典,为后面的计算做准备。在实际情况中,词典通常会很大,为了避免重复计算,可以将词典保存下来。

1
2
3
dictionary = corpora.Dictionary(texts)
dictionary.save('sample.dict')
print(dictionary)
Dictionary(10 unique tokens: ['不错', '你好', '好的', '笨蛋', '傻子']...)

有了前一步生成的词典,我们就可以将每个词语转化成一个索引,表示该词语在所有文档中出现的次序,然后每个文档就可以转化成一个索引的列表,比如这里的第一个文档表示为[(0, 1), (1, 1), (2, 1), (3, 1)],它的含义是在第一篇文档中,第0个词语(也就是“你好”)出现了一次,第一个词语(也就是“好的”)出现了一次,依次类推。

1
2
3
corpus = [dictionary.doc2bow(text) for text in texts]
corpora.MmCorpus.serialize('sample.mm', corpus)
pprint(corpus)
[[(0, 1), (1, 1), (2, 1), (3, 1)],
 [(3, 1), (4, 1), (5, 1), (6, 1)],
 [(0, 1), (1, 1), (2, 1), (6, 1)],
 [(0, 1), (2, 1), (7, 1), (8, 1), (9, 1)],
 [(4, 1), (5, 1), (9, 1)],
 [(0, 1), (2, 1), (6, 1), (7, 1)],
 [(6, 2), (8, 2)]]

然后我们可以调用gensim中的tf-idf模块来进一步对每个单词计算权重。关于tf-idf的计算方法,可以参考维基百科

1
2
3
4
from gensim import models, similarities
tf_idf = models.TfidfModel(corpus)
vec = [(0, 1), (5, 1), (7, 1)]
print(tf_idf[vec])
[(0, 0.3011997233053068), (5, 0.6742695034927825), (7, 0.6742695034927825)]

现在,我们可以开始计算两个文档的相似度了。这里我们采用了余弦相似度作为衡量指标,当然还有其他的方式,可以参考文本相似度计算方法研究综述

1
2
3
index = similarities.SparseMatrixSimilarity(tf_idf[corpus], num_features=10)
sims = index[tf_idf[vec]]
print(sims)
[0.08686648 0.37695488 0.10641449 0.43870124 0.38928968 0.63969857
 0.        ]

我们还可以将整个相似度矩阵打印出来,就可以看到前面我们所准备的七篇文档两两之间的相似度。

1
print(index[tf_idf[corpus]])

[[1.         0.3609379  0.71441036 0.13975498 0.         0.20378576
  0.        ]
 [0.3609379  1.0000001  0.08823138 0.         0.64554304 0.08823138
  0.10185669]
 [0.71441036 0.08823138 1.0000001  0.17120475 0.         0.37446705
  0.1440983 ]
 [0.13975498 0.         0.17120475 1.         0.31315398 0.60019135
  0.4952337 ]
 [0.         0.64554304 0.         0.31315398 0.99999994 0.
  0.        ]
 [0.20378576 0.08823138 0.37446705 0.60019135 0.         1.
  0.1440983 ]
 [0.         0.10185669 0.1440983  0.4952337  0.         0.1440983
  0.99999994]]

本文中的代码使用Jupyter Notebook编写,需要的朋友可以直接到github上查看。

如果您阅读本文后有所收获,不妨打赏两块钱,您小小的支持都是对作者莫大的鼓励!