用gensim doc2vec计算文本相似度

最近开始接触gensim库，之前训练word2vec用Mikolov的c版本程序，看了很久才把程序看明白，在gensim库中，word2vec和doc2vec只需要几个接口就可以实现，实在是方便。python，我越来越爱你了。

这个程序很简单，直接上程序了。

# coding:utf-8

import sys
import gensim
import sklearn
import numpy as np

from gensim.models.doc2vec import Doc2Vec, LabeledSentence

TaggededDocument = gensim.models.doc2vec.TaggedDocument

def get_datasest():
    with open("out/wangyi_title_cut.txt", 'r') as cf:
        docs = cf.readlines()
        print len(docs)

    x_train = []
    #y = np.concatenate(np.ones(len(docs)))
    for i, text in enumerate(docs):
        word_list = text.split(' ')
        l = len(word_list)
        word_list[l-1] = word_list[l-1].strip()
        document = TaggededDocument(word_list, tags=[i])
        x_train.append(document)

    return x_train

def getVecs(model, corpus, size):
    vecs = [np.array(model.docvecs[z.tags[0]].reshape(1, size)) for z in corpus]
    return np.concatenate(vecs)

def train(x_train, size=200, epoch_num=1):
    model_dm = Doc2Vec(x_train,min_count=1, window = 3, size = size, sample=1e-3, negative=5, workers=4)
    model_dm.train(x_train, total_examples=model_dm.corpus_count, epochs=70)
    model_dm.save('model/model_dm_wangyi')

    return model_dm

def test():
    model_dm = Doc2Vec.load("model/model_dm_wangyi")
    test_text = ['《', '舞林', '争霸' '》', '十强' '出炉', '复活', '舞者', '澳门', '踢馆']
    inferred_vector_dm = model_dm.infer_vector(test_text)
    print inferred_vector_dm
    sims = model_dm.docvecs.most_similar([inferred_vector_dm], topn=10)


    return sims

if __name__ == '__main__':
    x_train = get_datasest()
    model_dm = train(x_train)

    sims = test()
    for count, sim in sims:
        sentence = x_train[count]
        words = ''
        for word in sentence[0]:
            words = words + word + ' '
        print words, sim, len(sentence[0])

我用了网页的热门娱乐新闻标题作为训练语料，输出结果如下，不知道是迭代次数不够还是怎么的，相似度都很低，不过语料只有4w条，且都是各式各样的八卦，很相似的句子确实少。

MJ 环球 春晚 复活 全场 尖叫 林俊杰 再现 经典   0.320982992649 11
 《 舞林 》 神 剪辑 遭热议 导师 面临 最 痛苦 抉择   0.278693914413 13
 辛晓琪 《 舞林 》 初试 舞步 令 评委 方俊 不淡定   0.251948922873 12
 众星 助阵 舞林 盛典 陈志朋 曝 小虎队 或 计划 巡演   0.249549359083 12
 《 舞林 》 首现 同 性别 组合 金星 杨丽萍 冲突 再起   0.247241020203 13
 焦恩俊 三 小时 备战 《 舞林 》 深情 绅士 获 金星 高分   0.241078704596 14
 《 舞 出 》 撒 贝宁 跳 苦情 舞 复活 陈冲 再现 经典   0.239552676678 15
 吸血 女王 惊现 《 舞林 》 何琳 直呼 手脚冰凉   0.238774299622 11
 虎年 春晚 看点 揭秘 黄宏 复活 《 整容 》 一次 过关   0.223410069942 13
 郭书瑶 《 舞林 》 拼 性感 场上 与 歌迷 遥相呼应   0.218700557947 12