如何用Python训练自定义Word2Vec模型并将类描述转为单向量？

1. 准备并训练Word2Vec模型 #

先用你的语料库训练Word2Vec模型，这里用gensim库实现，这是NLP领域处理词向量的常用工具。

步骤与代码 #

• 先做文本预处理：分词、过滤停用词和无效字符（英文用nltk，中文可替换为jieba）
• 加载语料（支持文本文件或CSV）
• 训练并保存模型

from gensim.models import Word2Vec
import nltk
from nltk.tokenize import word_tokenize
from nltk.corpus import stopwords
import pandas as pd

# 初始化预处理工具（英文场景）
nltk.download('punkt')
nltk.download('stopwords')
stop_words = set(stopwords.words('english'))

def preprocess_text(text):
    # 分词+转小写+过滤停用词/非字母字符
    tokens = word_tokenize(text.lower())
    filtered_tokens = [token for token in tokens if token.isalpha() and token not in stop_words]
    return filtered_tokens

# 加载语料：二选一即可
# 场景1：文本文件（每行一段文本）
corpus = []
with open('corpus.txt', 'r', encoding='utf-8') as f:
    for line in f:
        corpus.append(preprocess_text(line.strip()))

# 场景2：CSV文件（假设文本列名为'text'）
# df_corpus = pd.read_csv('corpus.csv')
# corpus = df_corpus['text'].apply(preprocess_text).tolist()

# 训练Word2Vec模型
model = Word2Vec(
    sentences=corpus,
    vector_size=100,  # 向量维度，可按需调整
    window=5,         # 上下文窗口大小
    min_count=1,      # 忽略出现次数少于此值的词
    workers=4         # 并行训练线程数
)
# 保存模型方便后续调用
model.save('word2vec_model.model')

2. 将类描述转换为单个向量 #

核心逻辑是把描述中所有有效词的向量做平均（最易用的方案，也可按需换成加权平均），得到整个描述的代表向量。

步骤与代码 #

• 加载训练好的模型
• 对类描述做同样的预处理
• 提取词向量并计算平均值（跳过不在模型词汇表的词）
• 把类名和对应向量保存到CSV

from gensim.models import Word2Vec
import pandas as pd
import numpy as np

# 加载训练好的模型
model = Word2Vec.load('word2vec_model.model')
# 复用之前的预处理函数（中文场景替换为jieba分词逻辑）
def preprocess_text(text):
    tokens = word_tokenize(text.lower())
    filtered_tokens = [token for token in tokens if token.isalpha() and token not in stop_words]
    return filtered_tokens

# 加载类描述CSV（假设列名为'class'和'description'）
df_classes = pd.read_csv('class_descriptions.csv')

def get_description_vector(description):
    tokens = preprocess_text(description)
    word_vectors = []
    # 收集所有在模型词汇表中的词向量
    for token in tokens:
        if token in model.wv:
            word_vectors.append(model.wv[token])
    # 处理无匹配词的情况：返回全0向量
    if not word_vectors:
        return np.zeros(model.vector_size)
    # 计算向量平均值
    return np.mean(word_vectors, axis=0)

# 生成每个类的向量
df_classes['vector'] = df_classes['description'].apply(get_description_vector)

# 可选：将向量拆分为单独列，方便后续分析
vector_columns = [f'vec_{i}' for i in range(model.vector_size)]
df_classes[vector_columns] = pd.DataFrame(df_classes['vector'].tolist(), index=df_classes.index)

# 保存结果
df_classes.to_csv('class_vectors.csv', index=False)

3. 进阶优化（可选） #

• TF-IDF加权平均：先计算语料库的TF-IDF权重，用权重给词向量加权后再求平均，让重要词的贡献更大
• 改用Doc2Vec：如果需要直接训练文档级向量，可使用gensim的Doc2Vec模型，无需事后合并词向量
• 过滤低频词：预处理时过滤语料中出现次数极少的词，减少噪声干扰

内容的提问来源于stack exchange，提问作者Rupshali Dasgupta