预测IMDb电影平均评分的文本嵌入方法

数月前，某社交平台数据科学板块出现一篇题为"未获反馈即遭数据科学岗位拒绝"的帖子。发帖人公开了其针对某家庭作业的Colab笔记本，要求是为公开的IMDb数据集构建预测电影平均评分的模型。

IMDb数据库由某中心运营，允许用户以1-10分评分电影，最终平均分将显示在电影页面显要位置。例如《肖申克的救赎》目前以930万用户投票获得的9.3分位居榜首。

数据特性分析

IMDb非商业数据集包含三个核心文件：

title.basics.tsv.gz：包含影片类型、上映年份、时长、类型等基础元数据
title.ratings.tsv.gz：存储影片平均评分和投票数
title.principals.tsv.gz：记录演职人员信息

为确保评分稳定性，仅分析获得至少30票的影片（约24.2万部）。值得注意的是，这些数据集缺乏制作公司、剧情摘要等关键信息，为预测模型构建增加了难度。

嵌入生成技术

采用阿里巴巴NLP团队的gte-modernbert-base文本嵌入模型，该模型基于ModernBERT架构优化，对JSON格式输入具有出色理解能力。通过以下技术流程生成嵌入：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19


device = "cuda:0"
dataloader = torch.utils.data.DataLoader(docs, batch_size=32,
                                         shuffle=False,
                                         pin_memory=True,
                                         pin_memory_device=device)

dataset_embeddings = []
for batch in tqdm(dataloader, smoothing=0):
    tokenized_batch = tokenizer(
        batch, max_length=8192, padding=True, truncation=True, return_tensors="pt"
    ).to(device)

    with torch.no_grad():
        outputs = model(**tokenized_batch)
        embeddings = outputs.last_hidden_state[:, 0].detach().cpu()
    dataset_embeddings.append(embeddings)

dataset_embeddings = torch.cat(dataset_embeddings)
dataset_embeddings = F.normalize(dataset_embeddings, p=2, dim=1)

使用某云平台的L4 GPU，耗时21分钟完成所有24.2万部影片的嵌入生成，成本仅0.10美元。

模型比较实验

方法一：支持向量机（SVM）

使用cuML库的GPU加速SVM实现，通过网格搜索优化超参数，最终测试集MSE达到1.087

方法二：多层感知机（MLP）

构建包含6个隐藏层（每层256单元）的MLP网络，使用0.6 dropout防止过拟合，最佳测试MSE为1.074

方法三：从头训练LLM

训练500万参数的自定义ModernBERT模型，仅用10轮训练即达到1.026的测试MSE，但存在明显过拟合现象

技术结论

实验表明，尽管预训练嵌入模型包含丰富先验知识，但针对特定领域任务，从头训练的专用语言模型可能表现更佳。所有嵌入驱动的方法均优于传统特征工程方法，为类似结构化数据预测任务提供了新的技术思路。

完整实现代码已开源在GitHub仓库，包含数据预处理、嵌入生成和模型训练的完整流程。

基于文本嵌入的IMDb电影评分预测技术解析

本文探讨使用大型语言模型生成电影元数据的文本嵌入，并比较支持向量机、多层感知机和从头训练的LLM在预测IMDb平均评分任务上的表现，最终发现从头训练的LLM模型取得了最佳预测效果。