Doc2Vec简单例子

Doc2Vec基于Word2Vec，作用是把一个段落，或者一个句子，一篇文档转化为一个向量，下面是简单的例子：

#!/usr/bin/env python# -*- coding: utf-8 -*-# created by fhqplzj on 2017/06/30 上午11:34"""Doc2Vec简单例子"""import osimport randomimport timefrom collections import Counterimport gensimimport gensim.models.doc2vecimport gensim.utilsimport smart_openfrom six.moves import xrangetest_data_dir = os.path.join(gensim.__path__[0], 'test', 'test_data')lee_train_file = os.path.join(test_data_dir, 'lee_background.cor')lee_test_file = os.path.join(test_data_dir, 'lee.cor')def read_corpus(fname, tokens_only=False):    with smart_open.smart_open(fname, encoding='iso-8859-1') as f:        for i, line in enumerate(f):            if tokens_only:                yield gensim.utils.simple_preprocess(line)            else:                yield gensim.models.doc2vec.TaggedDocument(gensim.utils.simple_preprocess(line), [i])# 带标签的训练集train_corpus = list(read_corpus(lee_train_file))# 不带标签的测试集test_corpus = list(read_corpus(lee_test_file, tokens_only=True))model = gensim.models.doc2vec.Doc2Vec(size=100, min_count=2, iter=55)model.build_vocab(train_corpus)# 使用BLAS，训练时间3秒内，不使用BLAS，训练时间2分钟内。start_time = time.time()model.train(train_corpus, total_examples=model.corpus_count, epochs=model.iter)elasped_time = time.time() - start_timeprint elasped_time# 预测print model.infer_vector('only you can prevent forrest fires'.split())# ranks存储的是与文档最相似的文档的id号，包括自身# second_ranks存储的是除文档自身外，与之最相似的文档（默认相似度排名第二的文档），包括文档号和相似度ranks, second_ranks = [], []for doc_id in xrange(len(train_corpus)):    inferred_vector = model.infer_vector(train_corpus[doc_id].words)    sims = model.docvecs.most_similar([inferred_vector], topn=len(model.docvecs))    # 最相似的文档的序号    rank = [docid for docid, sim in sims].index(doc_id)    ranks.append(rank)    # 第二相似的文档的信息    second_ranks.append(sims[1])# 打印出统计信息，0代表自己与自己最相似（正确的case），1代表自己与其他的最相似（错误的case）print Counter(ranks)# 从训练集中随机选择一篇文档doc_id = random.randint(0, len(train_corpus))# 打印出文档内容print u'Document ({}): «{}»\n'.format(doc_id, ' '.join(train_corpus[doc_id].words))# Doc2Vec模型参数print u'SIMILAR/DISSIMILAR DOCS PER MODEL %s:\n' % model# 文档向量化inferred_vector = model.infer_vector(train_corpus[doc_id].words)# 对于所有文档，按照相似度从大到小排序sims = model.docvecs.most_similar([inferred_vector], topn=len(model.docvecs))# 打印出最相似的、中间的、最不相似的文档for label, index in [('MOST', 0), ('MEDIAN', len(sims) // 2), ('LEAST', len(sims) - 1)]:    print u'%s %s: «%s»\n' % (label, sims[index], ' '.join(train_corpus[sims[index][0]].words))# 从训练集中随机选择一篇文档doc_id = random.randint(0, len(train_corpus))# 打印出文档内容print u'Train Document ({}): «{}»\n'.format(doc_id, ' '.join(train_corpus[doc_id].words))# 除文档自身外，与之最相似的文档的参数，包括文档号和相似度sim_id = second_ranks[doc_id]# 打印出除文档自身外，最相似的文档的内容print u'Similar Document {}: «{}»\n'.format(sim_id, ' '.join(train_corpus[sim_id[0]].words))# 从测试集中随机选择一篇文档，在训练集中找出最相似的文档doc_id = random.randint(0, len(test_corpus))inferred_vector = model.infer_vector(test_corpus[doc_id])sims = model.docvecs.most_similar([inferred_vector], topn=len(model.docvecs))print 'Test Document ({}): «{}»\n'.format(doc_id, ' '.join(test_corpus[doc_id]))print u'SIMILAR/DISSIMILAR DOCS PER MODEL %s:\n' % modelfor label, index in [('MOST', 0), ('MEDIAN', len(sims) // 2), ('LEAST', len(sims) - 1)]:    print u'%s %s: «%s»\n' % (label, sims[index], ' '.join(train_corpus[sims[index][0]].words))