word2vec_zh.py

import logging
import os
import sys
import multiprocessing
import jieba
import re
import pickle

from gensim.models import Word2Vec
from gensim.models.word2vec import LineSentence

def remove_tag(text):
    return re.sub(r'<.*?>', '', text)

def preprocess(file, output, separator=' '):
    data = list()

    with open(file) as f:
        lines = f.readlines()
        for line in lines:
            words = jieba.cut(remove_tag(line.strip()))
            data.append(separator.join(words))

    with open('wiki.pkl', 'w') as f:
        pickle.dump(data, f)

    with open(output, 'w') as f:
        for line in data:
            f.write(line.encode('utf-8') + '\n')

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))

    # remove html tag and cut words
    preprocess('./wiki_chs', './wiki')

    # gensim
    model = Word2Vec(LineSentence('./wiki'), size=200, window=5,
            min_count=5, workers=multiprocessing.cpu_count())
    
    model.save('wiki.model')
    model.wv.save_word2vec_format('word2vec.model', binary=False)