Gensim训练中文词向量实战

引言

实现文本匹配模型时经常需要预训练好的中文词/字向量，今天通过gensim和中文维基百科数据训练一个中文字向量。

安装相关包

首先要安装所需的几个包：

zhconv 1.4.3
gensim 4.3.2

由于中文维基百科是繁体字，需要用zhconv转换为简体；gensim就是我们用来训练word2vec所用的库了。

处理数据

数据选择最新的中文维基百科语料，大概有2.6G左右。注意该链接保存的是最新的数据，所以过几天去下载可能不同。

原始是xml格式的数据，里面包含各种标签，我们需要对其进行处理。幸运地是，gensim为我们实现了维基百科语料的处理类WikiCorpus，直接调用即可：

import jieba
import zhconv
from gensim.corpora import WikiCorpus
from zhconv import convertdef preprocess(text: str, min=1,max=1,lower: bool = True
) -> list[str]:if lower:text = text.lower()# 转换为简体中文text = convert(text, "zh-cn")# 按字拆分return list(text)if __name__ == "__main__":zh_wiki = "zhwiki-latest-pages-articles.xml.bz2"print(f"Processing {zh_wiki}")wiki_corpus = WikiCorpus(zh_wiki, tokenizer_func=preprocess)texts = wiki_corpus.get_texts()# 保存处理好的数据WIKI_SEG_TXT = "wiki_seg_char.txt"generator = wiki_corpus.get_texts()with open(WIKI_SEG_TXT, "w", encoding="utf-8") as output:for texts_num, tokens in enumerate(generator):output.write(" ".join(tokens) + "\n")if (texts_num + 1) % 100000 == 0:print(f"已处理{texts_num}篇文章")

本文介绍的是按字拆分的字向量，若要训练词向量，还需要进行分词，导入相关包，修改15行代码：
return list(jieba.cut(text))

处理好之后保存到wiki_seg_char.txt文件中：

-rw-rw-r-- 1 greyfoss greyfoss  4.4G Nov 22 18:30 wiki_seg_char.txt

可以看到，处理完之后有4.4G，我们可以看下处理后的内容：

$ head wiki_seg_char.txt欧 几 里 得 ， 公 元 前 三 世 纪 的 古 希 腊 数 学 家 ， 现 在 被 认 为 是 几 何 之 父 ， 此 画 为 拉 斐 尔 的 作 品 《 雅 典 学 院 》 ' ' ' 数 学 ' ' ' ， 是 研 究 数 量 、 结 构 以 及 空 间 等 概 念 及 其 变 化 的 一 门 学 科 ， 属 于 形 式 科 学 的 一 种 。 数 学 利 用 抽 象 化 和 逻 辑 推 理 ， 从 计 数 、 计 算 、 量 度 、 对 物 体 形 状 及 运 动 的 观 察 发 展 而 成 。 数 学 家 们 拓 展 这 些 概 念 ， 以 公 式 化 新 的 猜 想 ， 以 及 从 选 定 的 公 理 及 定 义 出 发 ， 严 谨 地 推 导 出 一 些 定 理 。 基 础 数 学 的 知 识 与 运 用 是 生 活 中 不 可 或 缺 的 一 环 。 对 数 学 基 本 概 念 的 完 善 ， 早 在 古 埃 及 、 美 索 不 达 米 亚 及 古 印 度 历 史 上 的 古 代 数 学 文 本 便 可 观 见 ， 而 在 古 希 腊 那 里 有 更 为 严 谨 的 处 理 。 从 那 时 开 始 ， 数 学 的 发 展 便 持 续 不 断 地 小 幅 进 展 ， 至 1 6 世 纪 的 文 艺 复 兴 时 期 ， 因 为 新 的 科 学 发 现 和 数 学 革 新 两 者 的 交 互 ， 致 使 数 学 的 加 速 发 展 ， 直 至 今 日 。 数 学 并 成 为 许 多 国 家 及 地 区 的 教 育 中 的 一 部 分 。 数 学 在 许 多 领 域 都 有 应 用 ， 包 括 科 学 、 工 程 、 医 学 、 经 济 学 和 金 融 学 等 。 数 学 对 这 些 领 域 的 应 用 通 常 被 称 为 应 用 数 学 ， 有 时 亦 会 激 起 新 的 数 学 发 现 ， 并 导 致 全 新 学 科 的 发 展 ， 例 如 物 理 学 的 实 质 性 发 展 中 建 立 的 某 些 理 论 激 发 数 学 家 对 于 某 些 问 题 的 不 同 角 度 的 思 考 。 数 学 家 也 研 究 纯 粹 数 学 ， 就 是 数 学 本 身 的 实 质 性 内 容 ， 而 不 以 任 何 实 际 应 用 为 目 标 。 许 多 研 究 虽 然 以 纯 粹 数 学 开 始 ， 但 其 过 程 中 也 发 现 许 多 可 用 之 处 。

注意，并没有进行去停止词、标点符号等处理。我们直接拿上面的结果开始训练字向量模型。

训练

from gensim.models import word2vec
import multiprocessing
import logging# 回调函数，可以打印损失
from gensim.models.callbacks import CallbackAny2Vec# 定义早停异常
class StopEarlyException(Exception):passclass Callback(CallbackAny2Vec):"""Callback to print loss after each epoch"""def __init__(self, patient=3):self.epoch = 0self.losses = []self.patient = patientself.counter = 0self.best_loss = float("inf")self.pre_sum = 0def on_epoch_end(self, model):# 这个Loss是累计值last_loss = model.get_latest_training_loss()if self.epoch == 0:loss = last_lossself.losses.append(loss)logging.info(f"Loss after epoch {self.epoch}: {loss:4f}")else:loss = last_loss - self.pre_sumself.losses.append(loss)logging.info(f"Loss after epoch {self.epoch}: {loss:4f}")if loss < self.best_loss:self.best_loss = lossself.counter = 0else:self.counter += 1if self.counter >= self.patient:raise StopEarlyException()self.pre_sum = last_lossself.epoch += 1logging.basicConfig(format="%(levelname)s - %(asctime)s: %(message)s",datefmt="%H:%M:%S",level=logging.INFO,
)max_cpu_counts = multiprocessing.cpu_count()
embedding_dim = 300logging.info(f"Use {max_cpu_counts} workers to train Word2Vec (dim={embedding_dim})")WIKI_SEG_TXT = "wiki_seg_char.txt"
# 读取训练数据
sentences = word2vec.LineSentence(WIKI_SEG_TXT)logging.info("begin train")
# 定义模型
model = word2vec.Word2Vec(vector_size=embedding_dim, workers=max_cpu_counts, min_count=5
)
# 构建词表
model.build_vocab(sentences, progress_per=1000000)
logging.info("build vocab finish")
# 训练模型
try:model.train(sentences,total_examples=model.corpus_count,epochs=5,report_delay=1,compute_loss=True,callbacks=[Callback()],)
except StopEarlyException:logging.info("break from stop early")
# 保存模型
output_model = f"word2vec.zh.{embedding_dim}.char.model"
model.wv.save_word2vec_format(output_model)logging.info("train finished")

注意，gensim的word2vec有一个坑，如果你想打印它的损失，直接打印你会发现损失会一直上升，原因是它计算的是所有epoch损失的累计值，为了打印每个epoch的损失需要进行特殊处理。

几乎都是采用了默认参数，由于训练语料较多，模型采用CBOW算法，但没有深入的调参，发现训练5个epoch的效果还不错。训练更多比如30个epoch损失可以更低，但实际结合模型效果并不好，应该是过拟合了。

并且这里定义了一个早停策略，epoch数大的话可能会触发。

最后打印的日志为：

INFO - 16:41:08: Word2Vec lifecycle event {'msg': 'training on 6219869715 raw words (4893530805 effective words) took 5204.8s, 940204 effective words/s', 'datetime': '2023-12-02T16:41:08.082832', 'gensim': '4.3.2', 'python': '3.10.13 (main, Sep 11 2023, 13:44:35) [GCC 11.2.0]', 'platform': 'Linux-5.4.0-150-generic-x86_64-with-glibc2.27', 'event': 'train'}
INFO - 16:41:08: storing 19651x300 projection weights into word2vec.zh.300.char.model
INFO - 16:41:11: train finished

这种写法我们只需要拿它保存的word2vec.zh.300.char.model文件就够了，在字向量的情况下也不大：

-rw-rw-r-- 1 greyfoss greyfoss   65M Dec  2 16:41 word2vec.zh.300.char.model
-rw-rw-r-- 1 greyfoss greyfoss   27M Dec  2 17:53 word2vec.zh.300.char.model.tar.gz

加载 & 测试

from gensim.models import KeyedVectorsoutput_model = f"word2vec.zh.300.char.model"print(f"Loading {output_model}...")
model = KeyedVectors.load_word2vec_format(output_model)result = model.most_similar('狗')
for word in result:print(word)

$ python evaluate.py 
Loading word2vec.zh.300.char.model...
('猫', 0.7125084400177002)
('犬', 0.604468047618866)
('猪', 0.5916579365730286)
('兔', 0.5814022421836853)
('狼', 0.5506706833839417)
('遛', 0.4923962652683258)
('鸡', 0.48695406317710876)
('猴', 0.48581114411354065)
('屎', 0.4827899634838104)
('驴', 0.48213639855384827)