Doc2Vec简单例子

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

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# created by fhqplzj on 2017/06/30 上午11:34
"""
Doc2Vec简单例子
"""
import os
import random
import time
from collections import Counter

import gensim
import gensim.models.doc2vec
import gensim.utils
import smart_open
from six.moves import xrange

test_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_time
print 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' % model
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))