使用gensim+jieba生成Word Embeddings

在自然语言处理中，使用Word2Vec来生成Word Embeddings是非常常用的一种做法，而gensim这个库提供了生成W2V的接口。

import os
import jieba
from gensim.models import Word2Vec

# 迭代器，使用jieba将句子进行分词
class Sentences(object):# 这个类可以根据实际情况重写，我已经将所有的文章进行分句，并整合到了一个文件里面
    def __init__(self, dirname):
        self.dirname = dirname # 句子所在文件，没句句子占一行
        jieba.load_userdict("wordBase.txt") # 加载词库

    def __iter__(self):
        for fname in os.listdir(self.dirname):
            for line in open(os.path.join(self.dirname, fname)):
                yield list(jieba.cut(line))


def train_word2vec(folder_path, size=100):
    sentences = Sentences(folder_path) #生成分词后的句子，是一个二维数组

    # size是词向量长度
    # worker是线程数量，建议与物理线程数量一致
    # min_count是指出现次数小于一定程度，就忽略，0表示不忽略
    model = Word2Vec(sentences, size=size, workers=8, min_count=0)

    # 训练结束就将模型保存起来
    model.save("word2vec_model")

# 生成50维度的词向量模型
train_word2vec("folder_path",50)

# 测试训练好的词向量模型，使用model[keyWord]即可获取keyword这个词的词向量
model = Word2Vec.load("word2vec_model")
print(model[keyWord])