gensim 中文文本相似度计算-蒲公英云

gensim 中文文本相似度计算

# -*- coding: utf-8 -*-
# __jiahuiyu__
import jieba
import logging
from gensim import corpora, models, similarities
from collections import defaultdict
logging.basicConfig(format='%(asctime)s : %(levelname)s : %(message)s', level=logging.INFO)
def a_sub_b(a, b):
    ret = []
    for el in a:
        if el not in b:
            ret.append(el)
    return ret
punctions = [' ', '，', '。', '、', '？', '~', '！', '@', '#', '￥', '%', '……', '&', '*', '（', '）', '——', '-', '`',
             '；', '：']
# stop_list = set([word for word in open('e:/descfile/reprocess/stopword_ch.txt')])
line_seg = []
stop_list = [line.strip() for line in open('e:/descfile/reprocess/stopword_ch.txt').readlines()]
print stop_list
fr = open('e:/descfile/reprocess/documents_test.txt', 'r')
document = fr.readlines()
corpora_documents = []
for line in document:
    item_str = []
    item = (jieba.cut(line.strip(), cut_all=False))
    for i in list(item):
        # i = i.encode('utf8')
        item_str.append(i)
    item_str = a_sub_b(item_str, list(stop_list))
    item_str = a_sub_b(item_str, list(punctions))
    corpora_documents.append(item_str)
# print corpora_documents
frequency = defaultdict(int)
for sen in corpora_documents:
    for token in sen:
        frequency[token] += 1
# print frequency
texts = [[token for token in text if frequency[token] > 1] for text in corpora_documents]
# print texts
dictionary = corpora.Dictionary(texts)
dictionary.save('e:/descfile/reprocess/test1.dict')
corpus = [dictionary.doc2bow(text) for text in texts]
corpora.MmCorpus.serialize('e:/descfile/reprocess/test1.mm', corpus)
tfidf = models.TfidfModel(corpus)
corpus_tfidf = tfidf[corpus]
lsi = models.LsiModel(corpus_tfidf, id2word=dictionary, num_topics=5)
corpus_lsi = lsi[corpus_tfidf]
lsi.save('e:/descfile/reprocess/test1.lsi')
# lsi = models.LsiModel.load('e:/descfile/reprocess/ch_model.lsi')
index = similarities.MatrixSimilarity(lsi[corpus])
index.save('e:/descfile/reprocess/test1.index')
# index = similarities.MatrixSimilarity.load('e:/descfile/reprocess/ch_test.index')
# test for similarity
new_sen = "恶人多作怪"
new_vec = dictionary.doc2bow(jieba.cut(new_sen, cut_all=False))
# print new_vec
lsi_vec = lsi[new_vec]
# print lsi_vec
sims = index[lsi_vec]
print list(enumerate(sims))
# sorted
simsorted = sorted(enumerate(sims), key=lambda item: -item[1])
print simsorted
""" for sen in corpora_documents: for w in sen: if w not in stop_list and punctions: print w """