使用Word2Vec计算电影与高评分影片相似度遇KeyError求助

错误原因 #

1. 词汇不匹配触发KeyError：训练Word2Vec模型时，用的是转小写后的分词结果，但生成高评分电影embeddings时，直接对原描述分词（未转小写），导致像Two这种大写词汇不在模型的词汇表中（模型里只有小写的two），触发键错误。
2. 最后一行代码语法与逻辑双重错误：word2vec_model.wv.similarity()需要两个输入参数（两个词或两组向量），但代码仅传入整个df['description']，既不符合语法要求，也没实现“计算每个电影与高评分影片相似度”的核心逻辑。

修复后的完整代码 #

import pandas as pd
from gensim.models import Word2Vec
from nltk.tokenize import word_tokenize
import numpy as np
from gensim.matutils import cossim

# 初始化数据集
df = pd.DataFrame ({
    'description': [
        'Two imprisoned men bond over a number of years',
        'A family heads to an isolated hotel for the winter',
        'In a future where technology controls everything',
        'A young lion prince flees his kingdom only to learn the true meaning of responsibility',
        'A group of intergalactic criminals are forced to work together to stop a fanatical warrior'
    ],
    'ratings': [8.7, 9.3, 7.9, 8.5, 8.1]
})

# 统一分词预处理：转小写后分词，避免词汇不匹配
def tokenize_text(text):
    return word_tokenize(text.lower())

df['tokenized_description'] = df['description'].apply(tokenize_text)

# 训练Word2Vec模型
word2vec_model = Word2Vec(
    sentences=df['tokenized_description'],
    vector_size=100,
    window=5,
    min_count=1,
    workers=4
)

# 筛选高评分电影，用copy避免修改原数据的警告
threshold = df['ratings'].quantile(0.75)
highest_rated_movies = df[df['ratings'] >= threshold].copy()

# 定义函数：将分词文本转换为句向量（取词向量的平均值）
def get_sentence_vector(tokenized_text, model):
    # 过滤掉不在模型词汇表中的词（兼容通用场景）
    valid_words = [word for word in tokenized_text if word in model.wv.key_to_index]
    if not valid_words:
        return np.zeros(model.vector_size)
    return np.mean(model.wv[valid_words], axis=0)

# 计算高评分电影的句向量集合
highest_rated_movies['sentence_vec'] = highest_rated_movies['tokenized_description'].apply(
    lambda x: get_sentence_vector(x, word2vec_model)
)
high_rated_vecs = highest_rated_movies['sentence_vec'].tolist()

# 定义函数：计算单个句向量与高评分句向量的平均余弦相似度
def calculate_avg_similarity(sentence_vec, high_rated_vecs):
    similarities = [cossim(sentence_vec, vec) for vec in high_rated_vecs]
    return np.mean(similarities)

# 计算所有电影的句向量及相似度得分
df['sentence_vec'] = df['tokenized_description'].apply(
    lambda x: get_sentence_vector(x, word2vec_model)
)
df['similarity_score'] = df['sentence_vec'].apply(
    lambda x: calculate_avg_similarity(x, high_rated_vecs)
)

# 输出结果
print(df[['description', 'ratings', 'similarity_score']])

关键修复说明 #

• 统一文本处理标准：所有分词操作先转小写，确保词汇与训练时的模型词汇表完全匹配。
• 句向量转换逻辑：Word2Vec是词向量模型，需将整段描述转换为句向量（这里采用词向量平均值，也可根据需求替换为加权平均等方法）。
• 相似度计算逻辑：用余弦相似度计算每个电影句向量与所有高评分电影句向量的相似度，取平均值作为最终特征得分。
• 规避数据修改警告：对筛选出的高评分电影数据框使用.copy()，防止修改切片数据时触发SettingWithCopyWarning。

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