doc2vct算法实现

本篇文章主要是实现Python 自然语言处理包 gensim 中用于长文本向量建模的 doc2vec算法。

示例代码如下：

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import logging
import multiprocessing
import os.path
import sys

from gensim import utils
from gensim.models.doc2vec import TaggedDocument, Doc2Vec


class MyDocs(object):
    def __init__(self, filename):
        self.filename = filename

    def __iter__(self):
        try:
            for line in open(self.filename, 'rb'):
                pieces = utils.to_unicode(line).split()
                tag = pieces[0]
                words = pieces[1].split(',')
                yield TaggedDocument(words, [tag])
        except:
            logging.info('e')


if __name__ == '__main__':
    program = os.path.basename(sys.argv[0])
    logger = logging.getLogger(program)

    logging.basicConfig(format='%(asctime)s: %(levelname)s: %(message)s')
    logging.root.setLevel(level=logging.INFO)
    logger.info("running %s" % ' '.join(sys.argv))

    # check and process input arguments
    if len(sys.argv) < 4:
        sys.exit(1)
    inp, outp1, outp2 = sys.argv[1:4]

    documents = MyDocs(inp)
    #建立模型
    model = Doc2Vec(documents, size=200, window=5, min_count=20, workers=multiprocessing.cpu_count())

    # trim unneeded model memory = use(much) less RAM
    # model.init_sims(replace=True)
    #保存模型
    model.save(outp1)
    model.save_word2vec_format(outp2, binary=False)

说明一下：在MyDocs这个类中，我自定义了一个逐步读取文件的方法，因为在试验中我们发现，一个4个G 左右大小的文件在一个服务器内存为64G 的服务器上是不够用的，内存消耗非常之快，所以就看了一下源码中TaggedLineDocument和LineSentence这个类的源码，发现使用 yield这个关键字能很好的解决这个问题，便有了MyDocs类，实验结果表明，效果大大提高。有兴趣的同学可以看看 yield的用法。

#测试模型
# -*- coding: utf-8 -*-

import sys

from gensim.models import Doc2Vec

if __name__ == '__main__':
    # check and process input arguments
    if len(sys.argv) < 3:
        sys.exit(1)
    file, word = sys.argv[1:3]

    #加载模型
    docvecs = Doc2Vec.load(file).docvecs
    print(docvecs.most_similar(word))

更多内容可以查看官方文档