gensim训练Word2Vec词向量模型

# encoding:utf-8
import multiprocessing
import re
import jieba.analyse
import gensim
import jieba.analyse
from gensim.models import Word2Vec
import os
import logging
from sklearn.manifold import TSNE
import numpy as np
import matplotlib.pyplot as plt
import random
import time

# 用来正常显示中文标签
plt.rcParams['font.sans-serif'] = ['SimHei']
# 用来正常显示负号
plt.rcParams['axes.unicode_minus'] = False

# 设置日志级别 in (debug、info、warning、error、critical) 只显示级别以上的日志
logging.basicConfig(level=logging.DEBUG, format=' %(asctime)s - %(levelname)s - %(message)s')


# 文件预处理，把文件按照？！。；分割
# 目的：把文本分为n句话
def __preProcess__(textPath):
    lines = open(textPath, encoding="utf-8", mode="r").read().replace("\n", "")
    # 去除空白数据
    lines = re.sub(re.compile(r"\s+"), '', lines)
    # 去除特殊符号
    specialSymbol = re.compile(r"[（）\[\]【】{}‘’~！@#￥%…&*-+=·—、；：/》《”，“`!^*()$-=_+'\";:\\|<>,.?]")
    lines = re.sub(specialSymbol, "", lines)
    # 按照结束符号分割
    lines = re.split("[？！。；]", lines)
    # 去除两端空白 例如'   ' 或者空白字符串
    lines = [i.strip() for i in lines if i.strip()]
    return lines


# 分词
def __splitWord__(lines):
    allWords = []
    # 启动paddle模式
    jieba.enable_paddle()
    # 加载停用词
    stopwords = [i.strip() for i in open(r"data/stop-words.txt", encoding="utf-8", mode="r").readlines()]
    for id, value in enumerate(lines):
        # 输出进度
        if (id + 1) % 500 == 0:
            status = (id + 1) / len(lines)
            logging.info(f'分词进度：{"%.2f%%" % (status * 100)}')
        # jieba分词后去除停用词
        lineWords = [i for i in jieba.lcut(value, use_paddle=True) if i not in stopwords]
        # 去除两端空白 例如'   ' 或者空字符串
        lineWords = [i.strip() for i in lineWords if i.strip()]
        # 如果分词后数组为空 不加入
        if lineWords:
            allWords.append(lineWords)
    return allWords


# 训练模型
def trainModel(textPath, modelName):
    # 数据预处理
    lines = __preProcess__(textPath)
    logging.info(f'数据预处理完成,共{len(lines)}条句子')
    # 分词
    allWords = __splitWord__(lines)
    logging.info('数据分词完成')
    """
   参数
   min_count 词频少于min_count次数的单词会被丢弃掉, 默认值为5。
   vector_size 向量的维度
   workers 多cpu
    """
    # 训练模型
    model = Word2Vec(allWords, workers=multiprocessing.cpu_count())
    logging.info('训练模型处理完成')
    # 保存模型
    __savaModel__(model, modelName)
    logging.info(f'保存模型处理完成(./model/gensim-model_{modelName})')
    return model


# 保存模型（两种形式）
def __savaModel__(model, modelName):
    # 如果文件夹不存在则创建
    if os.path.exists("model") and os.path.isdir("model"):
        pass
    else:
        os.mkdir("model")
    # 保存模型
    model.save(f"./model/gensim-model_{modelName}")
    # 保存word2vec（每个词对应的词向量值）
    model.wv.save_word2vec_format(f"./model/gensim-word2vec_{modelName}")


# 加载模型
def loadModel(modelName, word2vec=False, binary=False):
    # 模型路径
    modelPath = os.path.join("model", modelName)
    # 模型的两种方式
    if word2vec:
        model = gensim.models.KeyedVectors.load_word2vec_format(modelPath, binary=binary)
    else:
        model = gensim.models.Word2Vec.load(modelPath)
    logging.info('加载模型处理完成')
    return model


# 更新模型（后续使用更多句子和新词汇继续训练它：）
def updateModel(modelName, textPath, modelNewName):
    # 模型路径
    modelPath = os.path.join("model", modelName)
    # 加载旧模型
    model = gensim.models.Word2Vec.load(modelPath)
    # 数据预处理
    lines = __preProcess__(textPath)
    logging.info('数据预处理完成')
    # 分词
    allWords = __splitWord__(lines)
    logging.info('数据分词完成')
    model.build_vocab(allWords, update=True)
    # 训练
    model.train(allWords, total_examples=model.corpus_count, epochs=model.epochs)
    # 模型更新后的名字
    newModelName = modelName + "_update_" + time.strftime("%Y%m%d_%H:%M:%S", time.localtime())
    __savaModel__(model, newModelName)
    return model


# 降维（二维）
def __reduce_dimensions__(model):
    num_dimensions = 2

    vectors = np.asarray(model.wv.vectors)
    labels = np.asarray(model.wv.index_to_key)

    # 使用 t-SNE降维
    tsne = TSNE(n_components=num_dimensions, random_state=0)
    vectors = tsne.fit_transform(vectors)

    x_vals = [v[0] for v in vectors]
    y_vals = [v[1] for v in vectors]
    return x_vals, y_vals, labels


# 画图
def __plot_with_matplotlib__(x_vals, y_vals, labels, pictureName):
    random.seed(0)

    plt.figure(figsize=(12, 12))
    plt.scatter(x_vals, y_vals)

    indices = list(range(len(labels)))
    selected_indices = random.sample(indices, 25)
    for i in selected_indices:
        plt.annotate(labels[i], (x_vals[i], y_vals[i]))
    # 保存svg
    plt.savefig(f"{pictureName}.svg", format="svg")
    plt.show()


# 模型可视化
def modelVisualization(model, pictureName):
    # 降维
    x_vals, y_vals, labels = __reduce_dimensions__(model)
    # 绘图
    __plot_with_matplotlib__(x_vals, y_vals, labels, pictureName)

# 训练模型
trainModel(r"data/斗罗大陆.txt","斗罗大陆")
model = loadModel("gensim-model_斗罗大陆")
# 输出相似
for i in ['唐三', '小舞', '戴沐白', '朱竹清', '奥斯卡', '宁荣荣', '马红俊']:
	print(f"和{i}最相似的人是{model.wv.most_similar(positive=i, topn=1)}")
# 可视化
modelVisualization(model, "斗罗大陆")

用斗罗大陆训练了一个词向量模型

奥斯卡和宁荣荣真爱无疑！