一次gensim的Word2Vec体验之旅

前言

最近做NLP的任务中经常性地用到了word2vec方法，索性写一篇文章来加深理解和使用。本文主要用到了gensim包的Word2Vec类，word2vec本质上是将单词转换成了计算机接收的向量，用较低维度的表示体现了单词之间的语义信息，关于word2vec原理的理解，这篇博客讲述地十分通俗易懂了，这里不多赘述，我们主要关注如何通过给定语料库获得词向量的表示。

训练语料

要利用Word2Vec训练词向量，需要有充足的语料数据，这里以wiki中文语料库为例。

wiki中文语料库原地址下载较慢（国外服务器），这里提供我上传的版本wiki中文语料库，若无法下载请评论留言。

利用gensim.corpora的WikiCorpus类，我们可以得到数据对象，用于具体内容的获取。
使用如下代码查看语料库的数据样例：

from gensim.corpora import WikiCorpus

wiki = WikiCorpus(inp, lemmatize=False, dictionary={})
	logging.info('transforming original file to txt ...')
	for text in wiki.get_texts():
		print(text)
		break

中文语料的数据样例如下图所示：

可以发现是繁体字数据，且包含一些英文单词，在将数据转换成txt的过程中去除这些英文单词，并将繁体字转换成简体字。

语料库数据处理

处理数据的代码如下所示：

import re
import os
import time
from gensim.corpora import WikiCorpus
from tqdm import tqdm
import zhconv

English_pattern = re.compile('[a-z]')  # 检测是否包含英文字母

def xml2txt(inp, outp, retransform=False):
    """
    处理原始文件-zhwiki-latest-pages-articles.xml.bz2成txt形式
    """
    if not os.path.exists(outp) or retransform:
        with open(outp, 'w', encoding='utf-8') as out:
            wiki = WikiCorpus(inp, lemmatize=False, dictionary={})
            print('Transforming original file to txt ...(it may takes about 30-40 mins)')
            for texts in tqdm(wiki.get_texts()):  # 每个text是一个列表
                # 去除text中的英文单词
                no_english_text = []
                for text in texts:
                    if not bool(English_pattern.search(text)):
                        # 未检测到英文字母
                        no_english_text.append(text)
                line_text = ' '.join(no_english_text) + '\n' # 拼接成一条字符串
                out.write(zhconv.convert(line_text, 'zh-cn'))  # 转换为简体输出
                    
            print('Transform finished.')

    else:
        print('Already Transformed')

查看wiki_zh.txt文件内容如下：

由于是中文数据，需要进行分词，采用jieba对txt文件内容分词，分词的代码如下：

def cut_words(inp, outp, recut =False):
    """
    对txt文件进行分词，保存到分词结果文件
    """
    if not os.path.exists(outp) or recut:
        print('Cutting words...')
        out = open(outp,'w',encoding='utf-8')
        with open(inp, 'r', encoding='utf-8') as f:
            for line in tqdm(f.readlines()):
                cut_result = jieba.lcut(line.strip())
                line_result = ' '.join(cut_result) + '\n'
                out.write(line_result)
        out.close()
    else:
        print('Already cutted.')

下面显示了分词后的结果:

Word2Vec模型训练

Word2Vec的训练很简单，调用Word2Vec类接收文本文件输入即可，其中有几个重要参数需要了解：

• LineSentence: 通常训练的语料数据是非常庞大的，我们无法一次加载到内存中，通过from gensim.models.word2vec import LineSentence导入使用，它接收输入文本，每次按行读取文件中的数据。
• size: 词向量的维度，一般像中文语料库这么大的至少用100以上的维度表示，我设置了200
• window: word2vec的窗口，考虑当前词左右window大小的词构建词向量，window的大小实际上决定了单词之间语义距离的远近，window越大，则离距离较远的两词之间的语义信息也会被考虑。
• min_count: 词频小于min_count的词不会被考虑，默认5
• workers: 同时工作进程数，用于并行加快速度，通常我们用multiprocessing.cpu_count()获取最大cpu数量来达到最佳速度。

下面的代码实现了对w2v模型的训练:

def train_w2v_model(inp, outp1, outp2, retrain=False):
    """
    输入是预处理好的训练数据，按行输入到Word2Vec中训练，最后保存模型和向量
    
    """ 
    if not os.path.exists(outp1) or not os.path.exists(outp2) or retrain:
        print('Training w2v model...')
        model = Word2Vec(LineSentence(inp), size=200, window=5, min_count=5, workers=multiprocessing.cpu_count())
        model.save(outp1)  # 用于加载
        model.wv.save_word2vec_format(outp2, binary=False)  # 用于查看
    else:
        print('w2v model alwarday trained.')

中文语料库十分庞大，训练的过程至少需要几个小时

加载训练好的模型后，最重要的就是word2vec的效果。下面的代码对模型进行了测试，测试结果如后：

def test_model():
    print('loading model...')
    model = Word2Vec.load('./w2v_model/wiki_zh_w2v_model.bin')
    print('苹果的词向量:{}'.format(model['苹果']))
    print('与苹果最接近单词:{}'.format(model.wv.most_similar('苹果')))
    print('苹果与葡萄相似度:{}，苹果与计算机相似度:{}'.format(model.wv.similarity('苹果','葡萄'),model.wv.similarity('苹果','计算机')))
    print('国王-男人+女人的词向量结果:{}'.format(model.most_similar(positive=['女人', '国王'], negative=['男人'])))

• 打印单词苹果的词向量
  

• 查找苹果相近的词
  

• 查看苹果和葡萄相似度 ，苹果与计算机相似度对比
  

• 查看词与词之间关系，例如国王+女人-男人的词向量最有可能是：
  

从结果来看，训练的模型较好地表示了单词的语义信息，这为进一步的自然语言处理工作打好坚实的基础。

完整代码见github

参考资料

• 使用 word2vec 训练wiki中英文语料库