Upload scripts/variance_select.py with huggingface_hub

scripts/variance_select.py

@@ -0,0 +1,312 @@
+"""
+Variance-based One-Shot Question Selection
+===========================================
+对 SFT 模型在 1533 道开放题上各跑 10 次推理，
+用精确字符串匹配判断对错，选出方差最大的题目。
+
+用法:
+  docker exec rl4phyx_env python3 /workspace/rl4phyx/RL4Phyx/SFT/variance_select.py
+
+输出:
+  - variance_results.json: 每道题的正确率和方差
+  - best_question_for_rlvr.json: 方差最大的题目信息
+  - rlvr_train.parquet: 转好的训练数据 (1题 × 128行)
+"""
+
+import json
+import re
+import os
+import torch
+import numpy as np
+from pathlib import Path
+
+# ============ 配置 ============
+MODEL_PATH = "/workspace/rl4phyx/RL4Phyx/SFT/checkpoints/sft_qwen25vl_3b/merged"
+TEST_FILE = "/workspace/rl4phyx/RL4Phyx/SFT/sft_test.jsonl"
+IMAGE_DIR = "/workspace/rl4phyx/RL4Phyx/SFT/images"
+OUTPUT_DIR = "/workspace/rl4phyx/RL4Phyx/SFT"
+
+NUM_RUNS = 10          # 每题推理次数
+MAX_NEW_TOKENS = 1024  # 最大生成长度 (缩短加速，只需提取boxed答案)
+TEMPERATURE = 0.7      # 采样温度
+BATCH_SIZE = 4         # 推理 batch size (根据显存调)
+RLVR_COPIES = 128      # one-shot 复制次数
+
+
+def extract_boxed(text):
+    """从模型输出中提取 \\boxed{} 内的答案"""
+    # 找最后一个 \boxed{}
+    matches = re.findall(r'\\boxed\{([^{}]*(?:\{[^{}]*\}[^{}]*)*)\}', text)
+    if matches:
+        return matches[-1].strip()
+    return None
+
+
+def normalize_answer(ans):
+    """轻度归一化：去空格、去末尾句号"""
+    if ans is None:
+        return None
+    ans = ans.strip()
+    ans = ans.rstrip('.')
+    # 去掉 \text{} 包裹
+    ans = re.sub(r'\\text\{([^}]*)\}', r'\1', ans)
+    # 去多余空格
+    ans = re.sub(r'\s+', ' ', ans)
+    return ans
+
+
+def exact_match(pred_answer, gt_answer):
+    """精确字符串匹配"""
+    if pred_answer is None:
+        return False
+    pred_norm = normalize_answer(pred_answer)
+    gt_norm = normalize_answer(gt_answer)
+    if pred_norm is None or gt_norm is None:
+        return False
+    return pred_norm == gt_norm
+
+
+def load_test_data(test_file):
+    """加载测试数据"""
+    data = []
+    with open(test_file, 'r', encoding='utf-8') as f:
+        for line in f:
+            r = json.loads(line.strip())
+            data.append(r)
+    print(f"Loaded {len(data)} test samples")
+    return data
+
+
+def run_inference(model, processor, test_data, run_id):
+    """对所有题目跑一次推理"""
+    print(f"\n{'='*60}")
+    print(f"  Run {run_id + 1}/{NUM_RUNS}")
+    print(f"{'='*60}")
+    
+    results = []
+    
+    for i, item in enumerate(test_data):
+        if i % 10 == 0:
+            print(f"  Processing {i}/{len(test_data)}...", flush=True)
+        
+        prompt_text = item['prompt']
+        image_path = os.path.join(IMAGE_DIR, f"{item['index']}.png")
+        
+        # 构建消息
+        messages = [{"role": "user", "content": []}]
+        
+        # 添加图片（如果存在）
+        if os.path.exists(image_path):
+            messages[0]["content"].append({
+                "type": "image",
+                "image": f"file://{image_path}"
+            })
+        
+        messages[0]["content"].append({
+            "type": "text",
+            "text": prompt_text
+        })
+        
+        # 处理输入
+        text = processor.apply_chat_template(
+            messages, tokenize=False, add_generation_prompt=True
+        )
+        
+        from qwen_vl_utils import process_vision_info
+        image_inputs, video_inputs = process_vision_info(messages)
+        
+        inputs = processor(
+            text=[text],
+            images=image_inputs,
+            videos=video_inputs,
+            padding=True,
+            return_tensors="pt"
+        ).to(model.device)
+        
+        # 生成
+        with torch.no_grad():
+            output_ids = model.generate(
+                **inputs,
+                max_new_tokens=MAX_NEW_TOKENS,
+                temperature=TEMPERATURE,
+                do_sample=True,
+                top_p=0.9,
+            )
+        
+        # 解码 (只取生成部分)
+        input_len = inputs['input_ids'].shape[1]
+        generated = output_ids[0][input_len:]
+        prediction = processor.decode(generated, skip_special_tokens=True)
+        
+        # 提取答案
+        pred_answer = extract_boxed(prediction)
+        gt_answer = item['ground_truth']
+        is_correct = exact_match(pred_answer, gt_answer)
+        
+        results.append({
+            'index': item['index'],
+            'pred_answer': pred_answer,
+            'gt_answer': gt_answer,
+            'correct': is_correct,
+        })
+    
+    correct_count = sum(1 for r in results if r['correct'])
+    print(f"  Run {run_id + 1} accuracy: {correct_count}/{len(results)} "
+          f"({100*correct_count/len(results):.1f}%)")
+    
+    return results
+
+
+def compute_variance(all_runs, test_data):
+    """计算每道题的正确率方差"""
+    n_questions = len(test_data)
+    stats = []
+    
+    for qi in range(n_questions):
+        # 收集这道题在 10 次 run 中的对错情况
+        correct_flags = [all_runs[run_id][qi]['correct'] for run_id in range(NUM_RUNS)]
+        n_correct = sum(correct_flags)
+        p = n_correct / NUM_RUNS  # 正确率
+        variance = p * (1 - p)     # 伯努利方差
+        
+        stats.append({
+            'index': test_data[qi]['index'],
+            'category': test_data[qi].get('category', ''),
+            'ground_truth': test_data[qi]['ground_truth'],
+            'n_correct': n_correct,
+            'accuracy': p,
+            'variance': variance,
+            'correct_flags': correct_flags,
+            'pred_answers': [all_runs[run_id][qi]['pred_answer'] for run_id in range(NUM_RUNS)],
+        })
+    
+    # 按方差降序排列
+    stats.sort(key=lambda x: x['variance'], reverse=True)
+    
+    return stats
+
+
+def convert_to_training_format(question_item, copies=RLVR_COPIES):
+    """将选中的题目转成 RLVR 训练 parquet 格式"""
+    import pandas as pd
+    
+    prompt_text = question_item['prompt']
+    image_path = f"{question_item['index']}.png"
+    
+    # 构建 veRL 格式的 prompt
+    prompt_messages = [{"role": "user", "content": prompt_text}]
+    
+    records = []
+    for _ in range(copies):
+        records.append({
+            'prompt': prompt_messages,
+            'answer': question_item['ground_truth'],
+            'image_path': image_path,
+            'data_source': 'deepscaler',
+            'category': question_item.get('category', 'Physics'),
+            'index': question_item['index'],
+        })
+    
+    df = pd.DataFrame(records)
+    return df
+
+
+def main():
+    print("=" * 60)
+    print("  Variance-based One-Shot Question Selection")
+    print("=" * 60)
+    
+    # 1. 加载模型
+    print("\n[1/4] Loading SFT model...")
+    from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
+    
+    import os
+    os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+    
+    model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+        MODEL_PATH,
+        torch_dtype=torch.bfloat16,
+    ).to("cuda")
+    processor = AutoProcessor.from_pretrained(MODEL_PATH)
+    model.eval()
+    print(f"  Model loaded: {sum(p.numel() for p in model.parameters())/1e6:.0f}M params")
+    
+    # 2. 加载测试数据
+    print("\n[2/4] Loading test data...")
+    test_data = load_test_data(TEST_FILE)
+    
+    # 3. 跑 10 次推理
+    print(f"\n[3/4] Running {NUM_RUNS} inference passes on {len(test_data)} questions...")
+    all_runs = []
+    for run_id in range(NUM_RUNS):
+        run_results = run_inference(model, processor, test_data, run_id)
+        all_runs.append(run_results)
+        
+        # 每次 run 后保存中间结果
+        interim_path = os.path.join(OUTPUT_DIR, f"variance_run_{run_id}.json")
+        with open(interim_path, 'w', encoding='utf-8') as f:
+            json.dump(run_results, f, ensure_ascii=False, indent=2)
+        print(f"  Saved interim results to {interim_path}")
+    
+    # 4. 计算方差并选题
+    print(f"\n[4/4] Computing variance and selecting best question...")
+    stats = compute_variance(all_runs, test_data)
+    
+    # 保存完整统计
+    stats_path = os.path.join(OUTPUT_DIR, "variance_results.json")
+    with open(stats_path, 'w', encoding='utf-8') as f:
+        json.dump(stats, f, ensure_ascii=False, indent=2)
+    print(f"  Saved all variance stats to {stats_path}")
+    
+    # 打印 Top 20 最高方差题目
+    print(f"\n{'='*60}")
+    print(f"  TOP 20 HIGHEST VARIANCE QUESTIONS")
+    print(f"{'='*60}")
+    for i, s in enumerate(stats[:20]):
+        print(f"  #{i+1}: idx={s['index']} | gt={s['ground_truth'][:30]:30s} | "
+              f"correct={s['n_correct']}/{NUM_RUNS} | var={s['variance']:.4f} | "
+              f"cat={s['category']}")
+        print(f"       preds: {s['pred_answers'][:5]}")
+    
+    # 选方差最大的题
+    best = stats[0]
+    print(f"\n{'='*60}")
+    print(f"  SELECTED QUESTION FOR ONE-SHOT RLVR")
+    print(f"{'='*60}")
+    print(f"  Index:       {best['index']}")
+    print(f"  Category:    {best['category']}")
+    print(f"  Ground Truth: {best['ground_truth']}")
+    print(f"  Accuracy:    {best['n_correct']}/{NUM_RUNS} ({best['accuracy']*100:.0f}%)")
+    print(f"  Variance:    {best['variance']:.4f}")
+    print(f"  Pred Answers: {best['pred_answers']}")
+    
+    # 保存选中题目
+    best_idx = int(best['index'])
+    best_item = None
+    for item in test_data:
+        if int(item['index']) == best_idx:
+            best_item = item
+            break
+    
+    best_path = os.path.join(OUTPUT_DIR, "best_question_for_rlvr.json")
+    with open(best_path, 'w', encoding='utf-8') as f:
+        json.dump({
+            'selected_question': best_item,
+            'stats': best,
+        }, f, ensure_ascii=False, indent=2)
+    print(f"  Saved best question to {best_path}")
+    
+    # 转成训练 parquet
+    if best_item:
+        df = convert_to_training_format(best_item)
+        parquet_path = os.path.join(OUTPUT_DIR, "rlvr_train.parquet")
+        df.to_parquet(parquet_path, index=False)
+        print(f"  Saved training parquet ({len(df)} rows) to {parquet_path}")
+    
+    print(f"\n{'='*60}")
+    print(f"  DONE!")
+    print(f"{'='*60}")
+
+
+if __name__ == "__main__":
+    main()