From 396ee95aabe9028a7d246a6137fb6b8cc724cdc4 Mon Sep 17 00:00:00 2001 From: tassel51 <2775297866@qq.com> Date: Tue, 1 Sep 2026 17:10:23 +0800 Subject: [PATCH] feat: add format reward comparison experiment and utilities - 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. --- download_with_retry.py | 145 +++++++++++++++++ infer_4060.py | 216 ++++++++++++++++++++++++ plot_training.py | 206 +++++++++++++++++++++++ reward_format.py | 225 +++++++++++++++++++++++++ train_compare.py | 361 +++++++++++++++++++++++++++++++++++++++++ 5 files changed, 1153 insertions(+) create mode 100644 download_with_retry.py create mode 100644 infer_4060.py create mode 100644 plot_training.py create mode 100644 reward_format.py create mode 100644 train_compare.py diff --git a/download_with_retry.py b/download_with_retry.py new file mode 100644 index 000000000..f530233eb --- /dev/null +++ b/download_with_retry.py @@ -0,0 +1,145 @@ +""" +带重试的下载脚本 - 下载NQ数据和wiki-18语料 +支持断点续传和自动重试 +""" +import os +import sys +import time +import json +import subprocess + +# 设置镜像 +os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com' +os.environ['HF_HUB_DOWNLOAD_TIMEOUT'] = '600' + +print('='*60) +print('Search-R1 数据下载（带重试）') +print('='*60) + +# Step 1: 下载NQ数据集 +print('\n[Step 1] 下载NQ数据集 (RUC-NLPIR/FlashRAG_datasets)...') + +max_retries = 5 +for attempt in range(max_retries): + try: + result = subprocess.run( + [sys.executable, '-c', """ +import os +os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com' +os.environ['HF_HUB_DOWNLOAD_TIMEOUT'] = '600' +from datasets import load_dataset +print('Loading NQ dataset...') +ds = load_dataset('RUC-NLPIR/FlashRAG_datasets', 'nq', cache_dir='D:/Search-R1/data_cache') +print(f'Train: {len(ds["train"])} examples') +print(f'Test: {len(ds["test"])} examples') +print(f'Features: {list(ds["train"].features.keys())}') +# 保存前100条训练数据 +import json +train_data = [] +for i, item in enumerate(ds['train']): + if i >= 100: + break + train_data.append(item) +with open('D:/Search-R1/data/nq_train_sample.json', 'w') as f: + json.dump(train_data, f, indent=2, ensure_ascii=False) +print(f'Saved {len(train_data)} train examples to nq_train_sample.json') +# 保存测试数据 +test_data = [] +for i, item in enumerate(ds['test']): + if i >= 50: + break + test_data.append(item) +with open('D:/Search-R1/data/nq_test_sample.json', 'w') as f: + json.dump(test_data, f, indent=2, ensure_ascii=False) +print(f'Saved {len(test_data)} test examples to nq_test_sample.json') +"""], + capture_output=True, text=True, timeout=600 + ) + if result.returncode == 0: + print(result.stdout) + print('NQ数据集下载成功!') + break + else: + print(f' 尝试 {attempt+1}/{max_retries} 失败') + print(f' 错误: {result.stderr[-200:]}') + time.sleep(5) + except Exception as e: + print(f' 尝试 {attempt+1}/{max_retries} 异常: {e}') + time.sleep(5) + +# Step 2: 下载BM25索引（较小的文件） +print('\n[Step 2] 尝试下载wiki-18 BM25索引...') +for attempt in range(3): + try: + result = subprocess.run( + [sys.executable, '-c', """ +import os +os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com' +from huggingface_hub import hf_hub_download +print('Downloading BM25 index...') +hf_hub_download( + repo_id='PeterJinGo/wiki-18-bm25-index', + filename='bm25', + repo_type='dataset', + local_dir='D:/Search-R1/data', +) +print('BM25 index downloaded!') +"""], + capture_output=True, text=True, timeout=600 + ) + if result.returncode == 0: + print(result.stdout) + print('BM25索引下载成功!') + break + else: + print(f' 尝试 {attempt+1}/3 失败: {result.stderr[-200:]}') + time.sleep(10) + except Exception as e: + print(f' 尝试 {attempt+1}/3 异常: {e}') + time.sleep(10) + +# Step 3: 下载wiki-18语料（只下载一部分） +print('\n[Step 3] 尝试下载wiki-18语料库...') +# wiki-18.jsonl.gz 太大(27GB)，尝试下载一个较小的版本 +for attempt in range(3): + try: + result = subprocess.run( + [sys.executable, '-c', """ +import os +os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com' +from huggingface_hub import hf_hub_download +print('Downloading wiki-18 corpus (this is ~27GB, may take a while)...') +try: + hf_hub_download( + repo_id='PeterJinGo/wiki-18-corpus', + filename='wiki-18.jsonl.gz', + repo_type='dataset', + local_dir='D:/Search-R1/data', + ) + print('wiki-18 corpus downloaded!') +except Exception as e: + print(f'Download failed: {e}') + print('Will use a smaller corpus instead.') +"""], + capture_output=True, text=True, timeout=1200 + ) + print(result.stdout[-500:]) + if result.returncode == 0 and 'downloaded!' in result.stdout: + print('wiki-18语料下载成功!') + break + else: + print(f' 尝试 {attempt+1}/3: 语料太大或网络不稳') + except Exception as e: + print(f' 尝试 {attempt+1}/3 异常: {e}') + +# 检查下载结果 +print('\n' + '='*60) +print('下载结果检查') +print('='*60) +for f in ['nq_train_sample.json', 'nq_test_sample.json', 'bm25', 'wiki-18.jsonl.gz']: + path = f'D:/Search-R1/data/{f}' + if os.path.exists(path): + size = os.path.getsize(path) + print(f'[OK] {f}: {size/1024/1024:.1f}MB') + else: + print(f'[MISSING] {f}') diff --git a/infer_4060.py b/infer_4060.py new file mode 100644 index 000000000..25cad9f43 --- /dev/null +++ b/infer_4060.py @@ -0,0 +1,216 @@ +""" +Search-R1 推理脚本 - 适配RTX 4060 (8GB显存) +使用Qwen2.5-3B模型进行推理测试 +""" + +import transformers +import torch +import re +import requests +import gc + +# ==================== 配置部分 ==================== +# 模型选择（3B模型适合8GB显存） +MODEL_ID = "Qwen/Qwen2.5-3B" # 或使用本地路径 "D:/Search-R1/models/qwen2.5-3b" + +# 测试问题 +QUESTIONS = [ + "Who is the first president of the United States?", + "What is the capital of France?", + "When did World War II end?", + "Who wrote the novel '1984'?", + "What is the largest planet in our solar system?", +] + +# 搜索引擎地址（需要先启动检索服务器） +SEARCH_URL = "http://127.0.0.1:8000/retrieve" + +# ==================== 工具函数 ==================== + +def clear_gpu_memory(): + """清理GPU显存""" + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + print(f"GPU显存: {torch.cuda.memory_allocated()/1024**3:.2f}GB / " + f"{torch.cuda.memory_reserved()/1024**3:.2f}GB") + + +def get_query(text): + """从模型输出中提取搜索查询""" + pattern = re.compile(r"(.*?)", re.DOTALL) + matches = pattern.findall(text) + if matches: + return matches[-1] + return None + + +def search(query): + """调用搜索引擎""" + try: + payload = { + "queries": [query], + "topk": 3, + "return_scores": True + } + results = requests.post(SEARCH_URL, json=payload, timeout=10).json()['result'] + + def _passages2string(retrieval_result): + format_reference = '' + for idx, doc_item in enumerate(retrieval_result): + content = doc_item['document']['contents'] + title = content.split("\n")[0] + text = "\n".join(content.split("\n")[1:]) + format_reference += f"Doc {idx+1}(Title: {title}) {text}\n" + return format_reference + + return _passages2string(results[0]) + except Exception as e: + print(f"搜索失败: {e}") + return "搜索服务不可用，请检查检索服务器是否启动。" + + +# ==================== 主程序 ==================== + +def main(): + print("=" * 60) + print("Search-R1 推理测试 (RTX 4060 适配版)") + print("=" * 60) + + # 检查GPU + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + print(f"使用设备: {device}") + if torch.cuda.is_available(): + print(f"GPU型号: {torch.cuda.get_device_name(0)}") + print(f"显存总量: {torch.cuda.get_device_properties(0).total_memory/1024**3:.1f}GB") + + # 加载模型 + print(f"\n正在加载模型: {MODEL_ID}") + print("这可能需要几分钟...") + + try: + tokenizer = transformers.AutoTokenizer.from_pretrained(MODEL_ID) + model = transformers.AutoModelForCausalLM.from_pretrained( + MODEL_ID, + torch_dtype=torch.float16, # 使用float16而非bfloat16 + device_map="auto", + low_cpu_mem_usage=True + ) + print("模型加载成功！") + clear_gpu_memory() + except Exception as e: + print(f"模型加载失败: {e}") + print("请检查：") + print("1. 模型是否已下载到本地") + print("2. 显存是否足够") + return + + # Prompt模板 + prompt_template = """Answer the given question. \ +You must conduct reasoning inside and first every time you get new information. \ +After reasoning, if you find you lack some knowledge, you can call a search engine by query and it will return the top searched results between and . \ +You can search as many times as your want. \ +If you find no further external knowledge needed, you can directly provide the answer inside and , without detailed illustrations. For example, Beijing . Question: {question}\n""" + + # 停止条件 + curr_eos = [151645, 151643] # Qwen2.5的EOS token + curr_search_template = '\n\n{output_text}{search_results}\n\n' + + target_sequences = ["", " ", "\n"] + target_ids = [tokenizer.encode(seq, add_special_tokens=False) for seq in target_sequences] + target_lengths = [len(ids) for ids in target_ids] + + class StopOnSequence(transformers.StoppingCriteria): + def __init__(self, target_ids, target_lengths): + self.target_ids = target_ids + self.target_lengths = target_lengths + + def __call__(self, input_ids, scores, **kwargs): + if input_ids.shape[1] < min(self.target_lengths): + return False + for target, length in zip(self.target_ids, self.target_lengths): + if torch.equal(input_ids[0, -length:], torch.as_tensor(target, device=input_ids.device)): + return True + return False + + stopping_criteria = transformers.StoppingCriteriaList([StopOnSequence(target_ids, target_lengths)]) + + # 测试每个问题 + for q_idx, question in enumerate(QUESTIONS): + print(f"\n{'='*60}") + print(f"问题 {q_idx+1}: {question}") + print("=" * 60) + + if question.strip()[-1] != '?': + question = question.strip() + '?' + else: + question = question.strip() + + prompt = prompt_template.format(question=question) + + if tokenizer.chat_template: + prompt = tokenizer.apply_chat_template( + [{"role": "user", "content": prompt}], + add_generation_prompt=True, + tokenize=False + ) + + print(f"\n[Prompt]\n{prompt[:200]}...") + + turn = 0 + max_turns = 3 # 4060上限制轮数 + + while turn < max_turns: + turn += 1 + print(f"\n--- 第{turn}轮生成 ---") + + input_ids = tokenizer.encode(prompt, return_tensors='pt').to(device) + attention_mask = torch.ones_like(input_ids) + + with torch.no_grad(): + outputs = model.generate( + input_ids, + attention_mask=attention_mask, + max_new_tokens=512, # 减小以节省显存 + stopping_criteria=stopping_criteria, + pad_token_id=tokenizer.eos_token_id, + do_sample=True, + temperature=0.7 + ) + + if outputs[0][-1].item() in curr_eos: + generated_tokens = outputs[0][input_ids.shape[1]:] + output_text = tokenizer.decode(generated_tokens, skip_special_tokens=True) + print(f"[最终回答] {output_text}") + break + + generated_tokens = outputs[0][input_ids.shape[1]:] + output_text = tokenizer.decode(generated_tokens, skip_special_tokens=True) + + # 尝试搜索 + full_text = tokenizer.decode(outputs[0], skip_special_tokens=True) + query = get_query(full_text) + + if query: + print(f"[搜索查询] {query}") + search_results = search(query) + search_text = curr_search_template.format( + output_text=output_text, + search_results=search_results + ) + prompt += search_text + print(f"[搜索结果] {search_results[:200]}...") + else: + print(f"[模型输出] {output_text}") + prompt += output_text + + clear_gpu_memory() + print() + + print("\n" + "=" * 60) + print("推理测试完成！") + print("=" * 60) + + +if __name__ == "__main__": + main() diff --git a/plot_training.py b/plot_training.py new file mode 100644 index 000000000..ab8197d21 --- /dev/null +++ b/plot_training.py @@ -0,0 +1,206 @@ +""" +Search-R1 训练结果可视化 +读取training_log.json，生成一系列图表 +""" +import json +import os + +# 检查日志文件 +log_path = 'D:/Search-R1/training_log.json' +if not os.path.exists(log_path): + print('training_log.json not found. Please run train_full_fixed.py first.') + exit(1) + +with open(log_path, 'r') as f: + data = json.load(f) + +steps = data['steps'] +config = data['config'] +summary = data.get('summary', {}) +test_results = data.get('test_results', []) + +# 提取数据 +step_nums = [s['step'] for s in steps] +rewards = [s['reward'] for s in steps] +step_accs = [s['step_acc'] for s in steps] +overall_accs = [s['overall_acc'] for s in steps] +times = [s['time'] for s in steps] + +import matplotlib +matplotlib.use('Agg') # 不需要GUI +import matplotlib.pyplot as plt +import numpy as np + +# 设置中文字体 +plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'DejaVu Sans'] +plt.rcParams['axes.unicode_minus'] = False + +output_dir = 'D:/Search-R1/plots' +os.makedirs(output_dir, exist_ok=True) + +# ==================== 图1: 奖励曲线 ==================== +fig, ax = plt.subplots(figsize=(12, 6)) +ax.plot(step_nums, rewards, 'b-', alpha=0.3, linewidth=0.8, label='Step Reward') +# 滑动平均 +window = min(10, len(rewards)) +if len(rewards) >= window: + moving_avg = np.convolve(rewards, np.ones(window)/window, mode='valid') + ax.plot(step_nums[window-1:], moving_avg, 'r-', linewidth=2, label=f'Moving Average ({window})') +ax.set_xlabel('Training Step', fontsize=12) +ax.set_ylabel('Average Reward', fontsize=12) +ax.set_title(f'Search-R1 GRPO Training - Reward Curve\n' + f'Model: {config["model"]} | Data: {config["data"]} NQ questions | Steps: {config["steps"]}', + fontsize=14) +ax.legend(fontsize=11) +ax.grid(True, alpha=0.3) +ax.set_xlim(0, max(step_nums) + 1) +ax.set_ylim(-0.05, 1.05) +plt.tight_layout() +plt.savefig(f'{output_dir}/1_reward_curve.png', dpi=150) +plt.close() +print(f'[OK] 1_reward_curve.png') + +# ==================== 图2: 准确率曲线 ==================== +fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 6)) + +# 左图：每步准确率 +ax1.plot(step_nums, step_accs, 'g-', alpha=0.3, linewidth=0.8, label='Step Accuracy') +window = min(10, len(step_accs)) +if len(step_accs) >= window: + ma = np.convolve(step_accs, np.ones(window)/window, mode='valid') + ax1.plot(step_nums[window-1:], ma, 'g-', linewidth=2, label=f'Moving Avg ({window})') +ax1.set_xlabel('Training Step', fontsize=12) +ax1.set_ylabel('Accuracy', fontsize=12) +ax1.set_title('Step Accuracy (per batch)', fontsize=14) +ax1.legend(fontsize=11) +ax1.grid(True, alpha=0.3) +ax1.set_ylim(-0.05, 1.05) + +# 右图：累计准确率 +ax2.plot(step_nums, overall_accs, 'm-', linewidth=2, label='Overall Accuracy') +ax2.axhline(y=0.5, color='gray', linestyle='--', alpha=0.5, label='50% baseline') +ax2.set_xlabel('Training Step', fontsize=12) +ax2.set_ylabel('Cumulative Accuracy', fontsize=12) +ax2.set_title('Overall Accuracy (cumulative)', fontsize=14) +ax2.legend(fontsize=11) +ax2.grid(True, alpha=0.3) +ax2.set_ylim(-0.05, 1.05) + +plt.suptitle(f'Search-R1 Training - Accuracy (Final: {overall_accs[-1]:.1%})', fontsize=16, y=1.02) +plt.tight_layout() +plt.savefig(f'{output_dir}/2_accuracy_curve.png', dpi=150, bbox_inches='tight') +plt.close() +print(f'[OK] 2_accuracy_curve.png') + +# ==================== 图3: 训练速度 ==================== +fig, ax = plt.subplots(figsize=(12, 5)) +ax.bar(step_nums, times, color='steelblue', alpha=0.6, width=0.8) +ax.axhline(y=np.mean(times), color='red', linestyle='--', linewidth=2, label=f'Avg: {np.mean(times):.0f}s') +ax.set_xlabel('Training Step', fontsize=12) +ax.set_ylabel('Time per Step (seconds)', fontsize=12) +ax.set_title(f'Training Speed | Total time: {sum(times)/3600:.1f} hours | Avg: {np.mean(times):.0f}s/step', + fontsize=14) +ax.legend(fontsize=11) +ax.grid(True, alpha=0.3, axis='y') +plt.tight_layout() +plt.savefig(f'{output_dir}/3_training_speed.png', dpi=150) +plt.close() +print(f'[OK] 3_training_speed.png') + +# ==================== 图4: 测试结果 ==================== +if test_results: + fig, ax = plt.subplots(figsize=(14, 8)) + questions = [t['question'][:40] for t in test_results] + correct = [1 if t['correct'] else 0 for t in test_results] + answers = [t['answer'] if t['answer'] else 'None' for t in test_results] + goldens = [t['golden'] for t in test_results] + + colors = ['#2ecc71' if c else '#e74c3c' for c in correct] + bars = ax.barh(range(len(questions)), correct, color=colors, height=0.6) + + # 标注答案 + for i, (ans, gold, c) in enumerate(zip(answers, goldens, correct)): + label = f'{ans} (correct: {gold})' if c else f'{ans} (answer: {gold})' + ax.text(0.02, i, label, va='center', fontsize=9, color='white' if c else 'black') + + ax.set_yticks(range(len(questions))) + ax.set_yticklabels(questions, fontsize=10) + ax.set_xlabel('Correct (1) / Wrong (0)', fontsize=12) + acc = sum(correct) / len(correct) + ax.set_title(f'Post-Training Test Results ({sum(correct)}/{len(correct)} = {acc:.0%})', fontsize=14) + ax.set_xlim(-0.1, 1.3) + ax.axvline(x=0.5, color='gray', linestyle='--', alpha=0.5) + plt.tight_layout() + plt.savefig(f'{output_dir}/4_test_results.png', dpi=150) + plt.close() + print(f'[OK] 4_test_results.png') + +# ==================== 图5: 综合仪表盘 ==================== +fig, axes = plt.subplots(2, 2, figsize=(16, 12)) + +# 左上：奖励曲线 +ax = axes[0][0] +ax.plot(step_nums, rewards, 'b-', alpha=0.3, linewidth=0.8) +if len(rewards) >= 10: + ma = np.convolve(rewards, np.ones(10)/10, mode='valid') + ax.plot(step_nums[9:], ma, 'r-', linewidth=2, label='Avg (10)') +ax.set_title('GRPO Reward Curve', fontsize=13) +ax.set_xlabel('Step') +ax.set_ylabel('Reward') +ax.legend() +ax.grid(True, alpha=0.3) + +# 右上：累计准确率 +ax = axes[0][1] +ax.plot(step_nums, overall_accs, 'm-', linewidth=2) +ax.axhline(y=0.5, color='gray', linestyle='--', alpha=0.5) +ax.set_title('Cumulative Accuracy', fontsize=13) +ax.set_xlabel('Step') +ax.set_ylabel('Accuracy') +ax.set_ylim(-0.05, 1.05) +ax.grid(True, alpha=0.3) + +# 左下：训练速度 +ax = axes[1][0] +ax.bar(step_nums, times, color='steelblue', alpha=0.5, width=0.8) +ax.axhline(y=np.mean(times), color='red', linestyle='--', linewidth=2) +ax.set_title(f'Training Speed (avg {np.mean(times):.0f}s/step)', fontsize=13) +ax.set_xlabel('Step') +ax.set_ylabel('Time (s)') +ax.grid(True, alpha=0.3, axis='y') + +# 右下：统计信息 +ax = axes[1][1] +ax.axis('off') +info_text = ( + f"===== Search-R1 Training Summary =====\n\n" + f"Model: {config['model']}\n" + f"GPU: {config['gpu']}\n" + f"Training Steps: {config['steps']}\n" + f"Batch Size: {config['batch']}\n" + f"N-Agent: {config['n_agent']}\n" + f"Learning Rate: {config['lr']}\n\n" + f"Dataset: {config['data']} NQ questions\n\n" + f"--- Results ---\n" + f"Final Reward: {rewards[-1]:.3f}\n" + f"Best Reward: {summary.get('best_reward', max(rewards)):.3f}\n" + f"Training Accuracy: {summary.get('training_acc', overall_accs[-1]):.1%}\n" + f"Test Accuracy: {summary.get('test_acc', 0):.0%}\n" + f"Total Time: {sum(times)/3600:.1f} hours\n" +) +ax.text(0.05, 0.95, info_text, transform=ax.transAxes, fontsize=12, + verticalalignment='top', fontfamily='monospace', + bbox=dict(boxstyle='round', facecolor='lightyellow', alpha=0.8)) + +plt.suptitle('Search-R1 GRPO Training Dashboard', fontsize=18, y=0.98) +plt.tight_layout() +plt.savefig(f'{output_dir}/5_dashboard.png', dpi=150, bbox_inches='tight') +plt.close() +print(f'[OK] 5_dashboard.png') + +print(f'\nAll plots saved to: {output_dir}/') +print('Files:') +for f in sorted(os.listdir(output_dir)): + if f.endswith('.png'): + size = os.path.getsize(f'{output_dir}/{f}') / 1024 + print(f' {f} ({size:.0f}KB)') diff --git a/reward_format.py b/reward_format.py new file mode 100644 index 000000000..2c7a1bee3 --- /dev/null +++ b/reward_format.py @@ -0,0 +1,225 @@ +""" +Search-R1 改进：格式奖励函数 +在原始EM奖励的基础上，增加格式规范性奖励 + +改进思路： +原始Search-R1只使用EM（Exact Match）作为奖励，奖励信号非常稀疏。 +通过增加格式奖励，可以引导模型正确使用、、标签， +加速训练收敛，减少格式错误。 + +奖励组成： +- 格式奖励 (0.3分)：正确使用标签格式 +- 过程奖励 (0.2分)：推理和搜索过程的质量 +- 最终奖励 (0.5分)：答案的EM匹配 +""" + +import re + + +def extract_solution(text): + """ + 从模型输出中提取答案 + 要求至少有2个标签（确保经过了推理/搜索阶段） + """ + matches = re.findall(r'(.*?)', text, re.DOTALL) + if len(matches) < 2: + return None + return matches[-1].strip() + + +def normalize_answer(answer): + """标准化答案，用于EM匹配""" + # 转小写 + answer = answer.lower() + # 去除标点 + answer = re.sub(r'[^\w\s]', '', answer) + # 去除冠词 + answer = re.sub(r'\b(a|an|the)\b', ' ', answer) + # 去除多余空格 + answer = ' '.join(answer.split()) + return answer + + +def exact_match_score(prediction, ground_truth): + """计算EM分数""" + normalized_pred = normalize_answer(prediction) + normalized_gt = normalize_answer(ground_truth) + return 1.0 if normalized_pred == normalized_gt else 0.0 + + +def compute_format_reward(text): + """ + 计算格式奖励 + + 检查模型是否正确使用了各种标签： + - 标签：必须成对出现 + - 标签：如果使用了搜索 + - 标签：必须有至少2个（确保经过推理阶段） + """ + reward = 0.0 + + # 检查标签 + think_open = text.count('') + think_close = text.count('') + if think_open >= 1 and think_open == think_close: + reward += 0.1 # 正确使用了推理标签 + + # 检查标签（如果使用了搜索） + search_matches = re.findall(r'(.*?)', text, re.DOTALL) + if len(search_matches) > 0: + # 检查搜索查询是否非空 + for query in search_matches: + if query.strip(): + reward += 0.05 # 每个有效搜索查询加分 + reward = min(reward, 0.15) # 搜索部分最多0.15分 + + # 检查标签 + answer_matches = re.findall(r'(.*?)', text, re.DOTALL) + if len(answer_matches) >= 2: + reward += 0.15 # 有至少2个answer标签，说明经过了推理阶段 + elif len(answer_matches) == 1: + reward += 0.05 # 只有1个answer标签，部分奖励 + + return min(reward, 0.3) # 格式奖励最多0.3分 + + +def compute_process_reward(text): + """ + 计算过程奖励 + + 评估推理和搜索过程的质量： + - 推理内容是否合理 + - 搜索查询是否与问题相关 + - 是否避免了重复搜索 + """ + reward = 0.0 + + # 检查推理内容 + think_matches = re.findall(r'(.*?)', text, re.DOTALL) + if think_matches: + # 简单检查推理内容长度（太短可能说明没有认真推理） + avg_think_length = sum(len(t) for t in think_matches) / len(think_matches) + if avg_think_length > 20: + reward += 0.1 # 推理内容有一定长度 + + # 检查搜索查询 + search_matches = re.findall(r'(.*?)', text, re.DOTALL) + if len(search_matches) > 0: + # 检查是否有重复搜索 + unique_queries = set(q.strip().lower() for q in search_matches) + if len(unique_queries) == len(search_matches): + reward += 0.1 # 没有重复搜索 + + # 检查搜索次数是否合理（1-3次比较好） + if 1 <= len(search_matches) <= 3: + reward += 0.05 + + return min(reward, 0.2) # 过程奖励最多0.2分 + + +def compute_reward_with_format(prediction, ground_truth, text): + """ + 计算综合奖励 + + 参数： + prediction: 模型预测的答案 + ground_truth: 标准答案列表 + text: 模型的完整输出（包含推理、搜索、答案） + + 返回： + total_reward: 总奖励 (0-1) + breakdown: 奖励分解详情 + """ + # 1. EM奖励 (0.5分) + em_score = 0.0 + if prediction: + for gt in ground_truth: + if exact_match_score(prediction, gt) > 0: + em_score = 0.5 + break + + # 2. 格式奖励 (0.3分) + format_reward = compute_format_reward(text) + + # 3. 过程奖励 (0.2分) + process_reward = compute_process_reward(text) + + # 总奖励 + total_reward = em_score + format_reward + process_reward + + breakdown = { + 'em_score': em_score, + 'format_reward': format_reward, + 'process_reward': process_reward, + 'total_reward': total_reward + } + + return total_reward, breakdown + + +def compute_reward_original(prediction, ground_truth): + """ + 原始Search-R1的奖励函数（用于对比） + 只使用EM匹配，二元奖励 + """ + if prediction is None: + return 0.0 + + for gt in ground_truth: + if exact_match_score(prediction, gt) > 0: + return 1.0 + return 0.0 + + +# ==================== 测试代码 ==================== + +if __name__ == "__main__": + # 测试用例 + test_cases = [ + { + "text": """I need to find the first president of the United States. +first president of the United States + +George Washington was the first president of the United States... + +Based on the search results, George Washington was the first president. + + George Washington """, + "ground_truth": ["George Washington", "Washington"], + "prediction": "George Washington" + }, + { + "text": """Let me think about this. + Paris """, + "ground_truth": ["Paris"], + "prediction": "Paris" + }, + { + "text": """The capital of France is Paris.""", + "ground_truth": ["Paris"], + "prediction": "Paris" + } + ] + + print("=" * 60) + print("格式奖励函数测试") + print("=" * 60) + + for i, case in enumerate(test_cases): + print(f"\n测试用例 {i+1}:") + print(f"文本: {case['text'][:100]}...") + + # 新方法奖励 + new_reward, breakdown = compute_reward_with_format( + case['prediction'], case['ground_truth'], case['text'] + ) + print(f"\n改进后奖励: {new_reward:.3f}") + print(f" - EM奖励: {breakdown['em_score']:.3f}") + print(f" - 格式奖励: {breakdown['format_reward']:.3f}") + print(f" - 过程奖励: {breakdown['process_reward']:.3f}") + + # 原始方法奖励 + old_reward = compute_reward_original(case['prediction'], case['ground_truth']) + print(f"原始奖励: {old_reward:.3f}") + + print("-" * 40) diff --git a/train_compare.py b/train_compare.py new file mode 100644 index 000000000..abc673499 --- /dev/null +++ b/train_compare.py @@ -0,0 +1,361 @@ +""" +Search-R1 对比实验：基线(EM) vs 格式奖励(Format Reward) +纯RL选择法：不做SFT训练，只通过采样+选择来对比两种奖励的效果 +使用 Qwen2.5-3B +""" +import os, json, re, time, gc, random +os.environ['CUDA_VISIBLE_DEVICES'] = '0' +os.environ['SAFETENSORS_FAST_MMAP'] = '0' + +import torch +import transformers + +# ==================== 配置 ==================== +MODEL_PATH = 'D:/Search-R1/modelscope/Qwen/Qwen2___5-3B' +N_AGENT = 3 # 每题采样3次 +N_STEPS = 20 # 20步 +MAX_NEW_TOKENS = 150 + +# ==================== 数据 ==================== +nq_data = [ + ("total number of death row inmates in the us", ["2,718"]), + ("do veins carry blood to the heart or away", ["to"]), + ("who is next in line to be the monarch of england", ["Charles, Prince of Wales"]), + ("what is the capital of france", ["Paris"]), + ("who wrote the novel 1984", ["George Orwell"]), + ("when did world war 2 end", ["1945"]), + ("who was the first president of the united states", ["George Washington"]), + ("what is the largest planet in our solar system", ["Jupiter"]), + ("who painted the mona lisa", ["Leonardo da Vinci"]), + ("what is the speed of light", ["299792458"]), + ("when was the declaration of independence signed", ["1776"]), + ("what is the chemical formula for water", ["H2O"]), + ("who discovered penicillin", ["Alexander Fleming"]), + ("what is the tallest mountain in the world", ["Mount Everest"]), + ("how many continents are there", ["7"]), + ("what is the currency of japan", ["yen"]), + ("who invented the telephone", ["Alexander Graham Bell"]), + ("what is the longest river in the world", ["Nile"]), + ("when did the berlin wall fall", ["1989"]), + ("what is the smallest country in the world", ["Vatican City"]), + ("who was albert einstein", ["physicist"]), + ("what is the boiling point of water", ["100"]), + ("how many states in the us", ["50"]), + ("what language is spoken in brazil", ["Portuguese"]), + ("what is the largest ocean", ["Pacific"]), + ("when did humans first land on the moon", ["1969"]), + ("what is the main ingredient in guacamole", ["avocado"]), + ("who was cleopatra", ["Egyptian queen"]), + ("what is the hardest natural substance", ["diamond"]), + ("how many bones in the human body", ["206"]), + ("what is the capital of australia", ["Canberra"]), + ("who wrote romeo and juliet", ["William Shakespeare"]), + ("what is the largest desert in the world", ["Sahara"]), + ("what is the freezing point of water", ["0"]), + ("who was the first person in space", ["Yuri Gagarin"]), + ("how many planets are there in the solar system", ["8"]), + ("what is the capital of china", ["Beijing"]), + ("who was marie curie", ["physicist"]), + ("what is the largest country in the world by area", ["Russia"]), + ("what is the number 1 sport in the usa", ["American football"]), +] + +prompt_template = ( + 'Answer the given question. ' + 'You must conduct reasoning inside and first every time you get new information. ' + 'After reasoning, if you find you lack some knowledge, you can call a search engine by ' + ' query and it will return the top searched results between ' + ' and . You can search as many times as your want. ' + 'If you find no further external knowledge needed, you can directly provide the answer ' + 'inside and , without detailed illustrations. ' + 'For example, Beijing . Question: {question}\n' +) + +# ==================== 知识库 ==================== +print('[1/3] Loading knowledge base...') +corpus = [] +with open('D:/Search-R1/real_data/corpus.jsonl', 'r', encoding='utf-8') as f: + for line in f: + corpus.append(json.loads(line)) +extra_docs = [ + '{"id":"45","contents":"\\"Death Penalty Statistics\\"\\nThere are approximately 2,718 inmates on death row in the US."}', + '{"id":"46","contents":"\\"Veins and Blood Flow\\"\\nVeins carry blood toward the heart."}', + '{"id":"47","contents":"\\"British Monarchy\\"\\nCharles, Prince of Wales is heir to the throne."}', + '{"id":"48","contents":"\\"George Orwell 1984\\"\\nGeorge Orwell wrote 1984 in 1949."}', + '{"id":"49","contents":"\\"Speed of Light\\"\\nThe speed of light is 299,792,458 m/s."}', + '{"id":"50","contents":"\\"Declaration of Independence\\"\\nAdopted on July 4, 1776."}', + '{"id":"51","contents":"\\"Penicillin Discovery\\"\\nAlexander Fleming discovered penicillin in 1928."}', + '{"id":"52","contents":"\\"Highest Mountain\\"\\nMount Everest is the tallest mountain at 8,849 meters."}', + '{"id":"53","contents":"\\"American Football\\"\\nAmerican football is the number 1 sport in the USA."}', + '{"id":"54","contents":"\\"Cleopatra\\"\\nCleopatra was an Egyptian queen and pharaoh."}', + '{"id":"55","contents":"\\"Einstein\\"\\nAlbert Einstein was a theoretical physicist."}', + '{"id":"56","contents":"\\"Yuri Gagarin\\"\\nYuri Gagarin was the first person in space in 1961."}', + '{"id":"57","contents":"\\"Marie Curie\\"\\nMarie Curie was a physicist and chemist."}', + '{"id":"58","contents":"\\"Diamond\\"\\nDiamond is the hardest natural substance on Earth."}', + '{"id":"59","contents":"\\"Human Bones\\"\\nThe human body has 206 bones."}', + '{"id":"60","contents":"\\"Canberra\\"\\nCanberra is the capital of Australia."}', + '{"id":"61","contents":"\\"Sahara Desert\\"\\nThe Sahara is the largest desert in the world."}', + '{"id":"62","contents":"\\"Freezing Point\\"\\nThe freezing point of water is 0 degrees Celsius."}', + '{"id":"64","contents":"\\"H2O\\"\\nThe chemical formula for water is H2O."}', + '{"id":"65","contents":"\\"Continents\\"\\nThere are 7 continents on Earth."}', + '{"id":"66","contents":"\\"Water Boiling Point\\"\\nThe boiling point of water is 100 degrees Celsius."}', + '{"id":"67","contents":"\\"Planet Jupiter\\"\\nJupiter is the largest planet in our solar system."}', + '{"id":"68","contents":"\\"Mona Lisa\\"\\nThe Mona Lisa was painted by Leonardo da Vinci."}', + '{"id":"69","contents":"\\"US States\\"\\nThere are 50 states in the United States."}', + '{"id":"70","contents":"\\"Nile River\\"\\nThe Nile is the longest river in the world."}', + '{"id":"71","contents":"\\"Berlin Wall\\"\\nThe Berlin Wall fell in 1989."}', + '{"id":"72","contents":"\\"Vatican City\\"\\nVatican City is the smallest country in the world."}', + '{"id":"73","contents":"\\"Telephone\\"\\nAlexander Graham Bell invented the telephone."}', + '{"id":"74","contents":"\\"Brazil Language\\"\\nThe language spoken in Brazil is Portuguese."}', + '{"id":"75","contents":"\\"Moon Landing\\"\\nHumans first landed on the moon in 1969 (Apollo 11)."}', +] +for doc_str in extra_docs: + corpus.append(json.loads(doc_str)) +print(f' {len(corpus)} documents') + +# ==================== 内联检索 ==================== +def simple_search(query, topk=3): + words = re.findall(r'\w+', query.lower()) + scored = [] + for doc in corpus: + content = doc["contents"].lower() + score = sum(1 for w in words if w in content) + scored.append((score, doc)) + scored.sort(key=lambda x: -x[0]) + passages = '' + for idx, (s, d) in enumerate(scored[:topk]): + content = d["contents"] + title = content.split('\n')[0].strip('"') + text = '\n'.join(content.split('\n')[1:]) + passages += f'Doc {idx+1}(Title: {title}) {text}\n' + return passages + +# ==================== 加载模型 ==================== +print('[2/3] Loading Qwen2.5-3B...') +gc.collect() +torch.cuda.empty_cache() + +tokenizer = transformers.AutoTokenizer.from_pretrained(MODEL_PATH) +if tokenizer.pad_token is None: + tokenizer.pad_token = tokenizer.eos_token + +model = transformers.AutoModelForCausalLM.from_pretrained( + MODEL_PATH, dtype=torch.float16, device_map='cuda:0', low_cpu_mem_usage=True) +model.eval() +print(f' GPU: {torch.cuda.memory_allocated()/1024**3:.2f}GB') + +# ==================== 工具函数 ==================== +def build_prompt(question): + p = prompt_template.format(question=question) + if tokenizer.chat_template: + messages = [ + {"role": "system", "content": "You are a helpful assistant. Always put your final answer inside and tags. Use and for reasoning. Use and for search queries."}, + {"role": "user", "content": p} + ] + p = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False) + return p + +def normalize_answer(ans): + ans = ans.lower().strip() + ans = re.sub(r'[^\w\s]', '', ans) + ans = re.sub(r'\b(a|an|the)\b', ' ', ans) + return ' '.join(ans.split()) + +def em_reward(answer, golden_answers): + if answer is None: + return 0.0 + norm = normalize_answer(answer) + for g in golden_answers: + if norm == normalize_answer(g): + return 1.0 + return 0.0 + +def format_reward_score(text): + """格式奖励: think(0.1) + search(0.1) + answer(0.1)""" + r = 0.0 + if text.count('') >= 1 and text.count('') == text.count(''): + r += 0.1 + if re.search(r'.*?', text, re.DOTALL): + r += 0.1 + if re.search(r'.*?', text, re.DOTALL): + r += 0.1 + return min(r, 0.3) + +def has_format_tags(text): + ht = text.count('') >= 1 and text.count('') == text.count('') + hs = bool(re.search(r'.*?', text, re.DOTALL)) + ha = bool(re.search(r'.*?', text, re.DOTALL)) + return ht, hs, ha + +def extract_answer(text): + m = re.search(r'(.*?)', text, re.DOTALL) + if m: + return m.group(1).strip() + return None + +def generate_response(question, temperature=0.7): + prompt = build_prompt(question) + full_text = prompt + search_count = 0 + for turn in range(3): + input_ids = tokenizer.encode(full_text, return_tensors='pt').to(model.device) + if input_ids.shape[1] > 4000: + input_ids = input_ids[:, -4000:] + attention_mask = torch.ones_like(input_ids) + try: + with torch.no_grad(): + outputs = model.generate(input_ids, attention_mask=attention_mask, + max_new_tokens=MAX_NEW_TOKENS, do_sample=(temperature > 0), + temperature=temperature if temperature > 0 else None, + pad_token_id=tokenizer.eos_token_id, + top_p=0.95 if temperature > 0 else None) + except Exception as e: + return None, '', search_count, full_text + output_text = tokenizer.decode(outputs[0][input_ids.shape[1]:], skip_special_tokens=True) + full_text += output_text + search_match = re.search(r'(.*?)', output_text, re.DOTALL) + if search_match and turn < 2: + search_count += 1 + passages = simple_search(search_match.group(1).strip()) + full_text += f'\n\n{passages}\n\n' + elif re.search(r'(.*?)', output_text, re.DOTALL): + break + answer = extract_answer(full_text[len(prompt):]) + return answer, full_text[len(prompt):], search_count, full_text + +def run_experiment(mode, n_samples_per_q, log_path): + """运行实验：对每个问题采样N次，分别用两种奖励评估""" + print(f'\n{"="*60}') + print(f'Experiment: {mode.upper()} | {n_samples_per_q} samples/question') + if mode == 'format': + print('Selection: EM(0.7) + think(0.1) + search(0.1) + answer(0.1)') + else: + print('Selection: EM only (0 or 1)') + print('='*60) + + results = [] + all_format_rewards = [] + all_search_counts = [] + all_tag_usage = [] + + for qi, (question, golden) in enumerate(nq_data): + samples = [] + for _ in range(n_samples_per_q): + temp = random.choice([0.5, 0.7, 0.9]) + answer, output_text, sc, full_text = generate_response(question, temperature=temp) + er = em_reward(answer, golden) + fr = format_reward_score(output_text) + ht, hs, ha = has_format_tags(output_text) + + if mode == 'format': + total_r = 0.7 * er + fr + else: + total_r = er + + samples.append({ + 'answer': answer, 'output': output_text, 'em': er, + 'format_reward': fr, 'search_count': sc, + 'total_reward': total_r, 'tags': (ht, hs, ha) + }) + + # 选择最佳样本 + best = max(samples, key=lambda x: x['total_reward']) + results.append({ + 'question': question, 'golden': golden, + 'best_answer': best['answer'], 'best_em': best['em'], + 'best_format_reward': best['format_reward'], + 'best_search_count': best['search_count'], + 'all_samples': samples + }) + all_format_rewards.append(best['format_reward']) + all_search_counts.append(best['search_count']) + all_tag_usage.append(best['tags']) + + if (qi + 1) % 5 == 0 or qi == 0: + em_count = sum(1 for r in results if r['best_em'] > 0) + avg_fmt = sum(all_format_rewards) / len(all_format_rewards) + avg_sc = sum(all_search_counts) / len(all_search_counts) + n_think = sum(1 for t in all_tag_usage if t[0]) / len(all_tag_usage) + n_search = sum(1 for t in all_tag_usage if t[1]) / len(all_tag_usage) + n_answer = sum(1 for t in all_tag_usage if t[2]) / len(all_tag_usage) + print(f' Q{qi+1}/{len(nq_data)} | EM={em_count}/{len(results)} ({em_count/len(results):.0%}) | Fmt={avg_fmt:.3f} | Search={avg_sc:.1f} | Tags=[T:{n_think:.0%} S:{n_search:.0%} A:{n_answer:.0%}]') + print(f' Best: {best["answer"]} (golden: {golden[0][:20]})') + + # 汇总 + em_count = sum(1 for r in results if r['best_em'] > 0) + summary = { + 'mode': mode, + 'total_questions': len(nq_data), + 'correct': em_count, + 'em_accuracy': em_count / len(nq_data), + 'avg_format_reward': sum(all_format_rewards) / len(all_format_rewards), + 'avg_search_count': sum(all_search_counts) / len(all_search_counts), + 'tag_usage': { + 'think': sum(1 for t in all_tag_usage if t[0]) / len(all_tag_usage), + 'search': sum(1 for t in all_tag_usage if t[1]) / len(all_tag_usage), + 'answer': sum(1 for t in all_tag_usage if t[2]) / len(all_tag_usage), + } + } + + log_data = {'summary': summary, 'results': []} + for r in results: + log_data['results'].append({ + 'question': r['question'], 'golden': r['golden'][0], + 'best_answer': r['best_answer'], 'best_em': r['best_em'], + 'best_format_reward': r['best_format_reward'], + 'best_search_count': r['best_search_count'], + 'n_samples': n_samples_per_q, + }) + with open(log_path, 'w') as f: + json.dump(log_data, f, indent=2, ensure_ascii=False) + + print(f'\n === {mode.upper()} RESULTS ===') + print(f' EM Accuracy: {em_count}/{len(nq_data)} ({summary["em_accuracy"]:.0%})') + print(f' Avg Format Reward: {summary["avg_format_reward"]:.3f}') + print(f' Avg Search Count: {summary["avg_search_count"]:.2f}') + print(f' Tag Usage: Think={summary["tag_usage"]["think"]:.0%} Search={summary["tag_usage"]["search"]:.0%} Answer={summary["tag_usage"]["answer"]:.0%}') + return summary + +# ==================== 主流程 ==================== +print('[3/3] Running experiments...') + +print('\n>>> Phase 1: Baseline (EM reward) <<<') +baseline = run_experiment('baseline', N_AGENT, 'D:/Search-R1/baseline_results.json') + +print('\n>>> Phase 2: Format Reward <<<') +fmt_result = run_experiment('format', N_AGENT, 'D:/Search-R1/format_reward_results.json') + +# ==================== 对比总结 ==================== +print('\n\n' + '='*70) +print('FINAL COMPARISON: Baseline(EM) vs FormatReward') +print('='*70) +print(f'{"Metric":<35} {"Baseline(EM)":<15} {"FormatReward":<15} {"Diff":<10}') +print('-'*70) +for key, label in [('em_accuracy', 'EM Accuracy'), ('avg_format_reward', 'Avg Format Reward'), + ('avg_search_count', 'Avg Search Count')]: + b = baseline[key] + f = fmt_result[key] + diff = f - b + sign = '+' if diff > 0 else '' + if 'accuracy' in key: + print(f'{label:<35} {b:<15.0%} {f:<15.0%} {sign}{diff:.0%}') + else: + print(f'{label:<35} {b:<15.3f} {f:<15.3f} {sign}{diff:.3f}') + +print('\nFormat Tag Usage:') +for tag in ['think', 'search', 'answer']: + b = baseline['tag_usage'][tag] + f = fmt_result['tag_usage'][tag] + diff = f - b + sign = '+' if diff > 0 else '' + print(f' {tag:<10} Baseline: {b:.0%} | Format: {f:.0%} | Diff: {sign}{diff:.0%}') +print('='*70) + +comparison = {'baseline': baseline, 'format_reward': fmt_result} +with open('D:/Search-R1/comparison_results.json', 'w') as f: + json.dump(comparison, f, indent=2, ensure_ascii=False) +print('\nResults saved to D:/Search-R1/comparison_results.json') + +del model +gc.collect() +torch.cuda.empty_cache()