如何将Word2Vec输出转换为适用于Sklearn的DataFrame格式？

以下是完整可复现的代码，直接实现将DataFrame文本列转为与CountVectorizer格式一致的Word2Vec均值向量，可直接传入Sklearn分类器：

import pandas as pd
from gensim.models import Word2Vec
from sklearn.feature_extraction.text import CountVectorizer
import numpy as np

# 1. 构造测试数据集
data = pd.DataFrame({
    'text': [
        'cat dog fish',
        'dog cat cat',
        'fish bird',
        'bird dog'
    ],
    'label': [0, 0, 1, 1]
})

# 2. 文本预处理：转为Word2Vec要求的分词列表格式
data['tokenized_text'] = data['text'].apply(lambda x: x.split())

# 3. 训练Word2Vec模型
w2v_model = Word2Vec(
    sentences=data['tokenized_text'],
    vector_size=10,  # 词向量维度，可根据任务调整为100-300
    window=2,
    min_count=1,  # 保留所有出现过的词，减少未登录词问题
    workers=4
)

# 4. 定义函数：计算单条文本的词向量均值
def get_mean_vector(tokens, model):
    # 过滤不在模型词表中的词
    valid_vectors = [model.wv[token] for token in tokens if token in model.wv]
    if not valid_vectors:
        # 若所有词都不在词表，返回全0向量保证格式统一
        return np.zeros(model.vector_size)
    return np.mean(valid_vectors, axis=0)

# 5. 生成所有文本的均值向量
data['w2v_mean_vector'] = data['tokenized_text'].apply(lambda x: get_mean_vector(x, w2v_model))

# 6. 转为Sklearn兼容的二维数组格式（与CountVectorizer输出一致）
w2v_features = np.vstack(data['w2v_mean_vector'].values)

# 对比CountVectorizer的输出格式
count_vec = CountVectorizer()
count_features = count_vec.fit_transform(data['text']).toarray()

# 验证两种特征的格式一致性（均为[样本数, 特征数]）
print("Word2Vec均值特征形状:", w2v_features.shape)
print("CountVectorizer特征形状:", count_features.shape)

关键细节说明 #

• 预处理要求：Word2Vec仅接受分词后的列表集合作为输入，英文直接用split()，中文需用分词工具（如jieba）拆分，替换为data['tokenized_text'] = data['text'].apply(lambda x: jieba.lcut(x))即可。
• 模型参数调整：min_count=1确保所有出现过的词都被纳入词表，避免因词频过低导致的未登录词错误；vector_size控制最终特征维度，通常设为100-300，可根据任务规模调整。
• 均值向量生成：必须过滤不在模型词表中的词，防止抛出KeyError；若整行无有效词向量，返回全0向量，确保每个样本都能生成符合格式的特征值。
• 格式适配Sklearn：用np.vstack将每行的一维向量堆叠为二维数组，与CountVectorizer输出的toarray()格式完全一致，可直接传入任意Sklearn分类器（如LogisticRegression、RandomForestClassifier）。

扩展优化 #

如果需要更好的效果，可使用预训练的Word2Vec模型（基于大规模语料训练），替换训练步骤为加载预训练模型即可：

# 加载本地预训练模型（需提前下载或训练）
pretrained_w2v = Word2Vec.load("pretrained_w2v.model")
# 直接用预训练模型生成均值向量
data['w2v_mean_vector'] = data['tokenized_text'].apply(lambda x: get_mean_vector(x, pretrained_w2v))

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