feat: format reward experiment + RTX 4060 support - #190
Open
Tassel51 wants to merge 1 commit into
Open
Conversation
- Add train_compare.py: Format reward vs baseline comparison experiment - Pure sampling-based comparison (no SFT training needed) - Works on RTX 4060 with 8GB VRAM - Tests EM accuracy and format compliance - Add reward_format.py: Format reward functions - think/search/answer tag detection - Combined reward scoring - Add infer_4060.py: Simplified inference for RTX 4060 - Inline search (no threading issues) - Chat template support - Add plot_training.py: Training visualization - Add download_with_retry.py: Download utility with retry logic Tested on Qwen2.5-3B with 40 NQ questions.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
大三学生复现 Search-R1 的时候做的一些小改进,分享给大家~
改了啥
1. 格式奖励对比实验 (train_compare.py)
复现的时候发现原始方案只有 EM 奖励,中间过程完全没有反馈,就加了个格式奖励试试。
奖励公式:EM(0.7) + think(0.1) + search(0.1) + answer(0.1)
实测结果(Qwen2.5-3B, 40道NQ题):
主要是想验证格式奖励能不能让模型更规范地用标签,没想到准确率也涨了一点哈哈。
2. RTX 4060 适配 (infer_4060.py)
我显卡是 4060 8GB 的,原版推理脚本跑起来有问题,就改了一个简化版:
3. 其他小工具
reward_format.py: 格式奖励函数plot_training.py: 训练曲线可视化download_with_retry.py: 带重试的下载脚本(网不好时候用的)测试环境
代码都在
train_compare.py里,直接python train_compare.py就能跑,不需要额外配置。如果有什么问题欢迎提 issue~