vinnylarouge/jevlike

来源：https://github.com/vinnylarouge/jevlike

Jevlike

训练一个小型模型，用于从动态变化的文本选项列表中进行选择。
Jev-like 模型接收一段文本和包含 N 个选项的列表，并为每个选项返回一个概率值。该过程通过一次前向传播完成，而非逐词生成答案。
Jev（https://typesafe.ai/blog/introducing-system-one-models-and-jev）是 TypeSafe 公司针对此类任务的商业模型，但其设计尚未公开。
本仓库提供了一个独立的基础模型，具备相同的输入输出结构。

演示

同一选项注意力头可对图像片段中的控制器按钮进行评分。
这段10秒影片 连接了两个选定的5秒片段：在七按钮Doom控制器上进行实时 deadly_corridor 战斗，随后是一个国际象棋控制器通过五个按键实现走位和行棋。
图表展示了每次决策使用的张量。
Doom片段来自提供的联合检查点，该检查点在录制的十局游戏中平均获得0.60次击杀和-97.50奖励。
国际象棋片段来自更强的专用象棋检查点，该检查点在50场采样对局中（对阵随机移动者）取得4胜46平0负，但对阵Stockfish等级0时为0胜2平48负。
片段选择基于活跃度，不代表常规表现或能力水平。

安装游戏扩展组件，并使用已发布的联合检查点录制新的640×480分辨率Doom运行轨迹：

uv pip install -e '.[games]'  
python examples/doom/play.py examples/checkpoints/joint-imitation.pt --episodes 10 --game-seconds 35.3 --device cpu --capture-resolution 640x480 --output runs/doom.mp4 --trace runs/doom-trace.json  

使用相同视觉布局渲染轨迹（生成静音影片，因作者拥有的音轨未包含在仓库中）：

(cd examples/film && npm install && npx playwright install chromium)  
examples/film/make-film.sh runs/doom-trace.json runs/doom-film.mp4 10  

发布内容包括Doom示例、国际象棋示例、单游戏检查点以及共享的12选项检查点。
两款游戏均使用 jevlike.vision 中的视觉评分器，示例中未重复部署模型副本。

架构

每个选项生成一个查询向量（表示文本的数字序列），该向量为上下文令牌分配注意力权重，形成对应选项的上下文向量。
共享点积运算将每个选项与上下文向量组合为单一评分，Softmax函数跨选项层将评分转换为概率总和。

默认编码器从零开始学习字节嵌入（编码器负责将文本转换为向量）。
可选路径使用冻结的预训练编码器，其权重在小型评分器学习过程中保持不变。

数据格式

每行一个JSON对象：

{"context":"客户需要退款。","options":["退款","销售","技术支持"],"label":0}  

label 为正确选项的从零起始索引。
每行可包含不同数量的选项（最少两个）。

快速开始

在仓库根目录执行以下命令（生成本地合成数据、训练模型、评估并预测）：

uv venv  
source .venv/bin/activate  
uv pip install -e '.[dev]'  
jevlike-data synthetic --output data/synthetic  
jevlike-train data/synthetic/train.jsonl --validation data/synthetic/validation.jsonl --output runs/synthetic.pt  
jevlike-eval runs/synthetic.pt data/synthetic/test.jsonl  
jevlike-predict runs/synthetic.pt --context "选择精确徽章琥珀獾。徽章：琥珀獾" --option "青灰鹤" --option "琥珀獾" --option "金鹭"  

评估结果包括：

• Top-1准确率（首选正确答案的比例）
• Top-3准确率（正确答案位于前三名的比例）
• 校准误差（置信度与观测准确率的差异）
• 打乱上下文对照组（将选项与错误上下文配对）
  模型性能应显著优于对照组。

使用自定义数据

1. 导出训练集、验证集和测试集JSONL文件（格式如上）
2. 每行应包含模型预测时可能遇到的所有选项
3. 相关联记录应保持同分布（如客户/页面）
4. 使用 jevlike-train 训练
5. 使用 jevlike-eval 评估保留测试集

默认编码器截断上下文至192字节、选项至32字节，可通过 --context-tokens / --option-tokens 调整。
训练支持CPU、Apple MPS（Mac GPU）及NVIDIA CUDA。

使用冻结预训练编码器

安装可选依赖后指定兼容的Hugging Face编码器：

uv pip install -e '.[transformers]'  
jevlike-train data/synthetic/train.jsonl --validation data/synthetic/validation.jsonl --output runs/qwen-head.pt --encoder hf --hf-model Qwen/Qwen2.5-0.5B --rank 256 --batch-size 8  

检查点仅保存评分器头和编码器名称（不复制冻结权重），加载时需访问相同模型。
--rank 参数控制评分器头宽度（影响可训练参数量和内存占用）。

Wikispeedia示例

scripts/get_wikispeedia.sh 下载SNAP公开数据集并构建下一次点击JSONL文件：

scripts/get_wikispeedia.sh  
jevlike-train data/wikispeedia/jsonl/train.jsonl --validation data/wikispeedia/jsonl/validation.jsonl --output runs/wikispeedia.pt  

引用：Robert West, Jure Leskovec. Human Wayfinding in Information Networks, WWW 2012。
数据集条款详见SNAP页面（https://snap.stanford.edu/data/wikispeedia.html）。

预期表现

• 合成菜单实验：一次评分达到约98%准确率
• Wikispeedia数据（目标分离）：冻结Qwen2.5-0.5B编码器+评分器达26%（对照组约8%）
• 从零训练小模型（40,000次点击）：达29%
• 八选项时，单次前向传播速度约为强制生成400标记的小型解码器的100倍

注：以上为实验数据，未声称达到Jev同等质量或复现TypeSafe私有方法。

局限性

• 本项目为研究基础框架，非Jev副本
• 性能依赖数据质量、数据划分及编码器选择
• 字节编码器成本低但语义理解能力有限
• 预训练路径需下载大模型，内存需求较高
• 单次评分需预先获取完整选项列表
• 速度对比基于本地小型解码器，非商业模型

许可证

代码遵循MIT许可证。下载数据集和预训练模型需遵守各自条款。