gensim计算文档相似度

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

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

from gensim import corpora
from pprint import pprint
documents = ["你好 好的 不错 笨蛋",
             "笨蛋 傻瓜 傻子 哈哈",
             "好的 不错 你好 哈哈",
             "有趣 可以 好的 不错 还行",
             "傻瓜 傻子 二货 还行",
             "可以 好的 不错 哈哈",
             "有趣 有趣 哈哈 哈哈"]
texts = [[word for word in document.split()] for document in documents]
pprint(texts)

[['你好', '好的', '不错', '笨蛋'],
 ['笨蛋', '傻瓜', '傻子', '哈哈'],
 ['好的', '不错', '你好', '哈哈'],
 ['有趣', '可以', '好的', '不错', '还行'],
 ['傻瓜', '傻子', '二货', '还行'],
 ['可以', '好的', '不错', '哈哈'],
 ['有趣', '有趣', '哈哈', '哈哈']]

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

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的计算方法，可以参考维基百科

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上查看。