"""PR-v1: Learned Repair Pilot on ZeroGPU (RTX Pro 6000, 48GB) Graft Planner v0.1 → Projection/Repair v1 Pre-registered: 100 steps, lr=2e-5, fp16, no sweeps.""" import spaces import gradio as gr import torch import json import re import os import gc import time from transformers import AutoModelForCausalLM, AutoTokenizer from peft import PeftModel from safetensors import safe_open from safetensors.torch import save_file from torch.optim import AdamW DATA_DIR = "/data" # ---- Module-level: load data (no GPU needed) ---- with open(f"{DATA_DIR}/data/few_shot_v2.json") as f: few_shot_data = json.load(f) shots = [] for ptype, data in few_shot_data.items(): shots.append(f"Problem: {data['prompt']}\n\nSolution:\n{data['solution']}\nFINAL_ANSWER: {data['expected_answer']}") SHOTS = '\n\n---\n\n'.join(shots) with open(f"{DATA_DIR}/data/mathcode_compose_v2.jsonl") as f: all_probs = [json.loads(line) for line in f] def build_prompt(p): return f"{SHOTS}\n\n---\n\nProblem: {p['prompt']}\n\nSolve step by step, then write your answer.\nFINAL_ANSWER:" def build_gpv2(source_dir): """Build GPV-2 adapter: L30-32 MLP only.""" tensors = {} with safe_open(f"{source_dir}/adapter_model.safetensors", framework='pt', device='cpu') as sf: for k in sf.keys(): tensors[k] = sf.get_tensor(k).clone() with open(f"{source_dir}/adapter_config.json") as f: cfg = json.load(f) gpv2 = {} kept = 0 for k, t in tensors.items(): parts = k.split('.') layer = int(parts[4]) mod = parts[5] if layer in {30, 31, 32} and mod == 'mlp': gpv2[k] = t.clone() kept += 1 else: gpv2[k] = torch.zeros_like(t) os.makedirs('/tmp/gpv2', exist_ok=True) save_file(gpv2, '/tmp/gpv2/adapter_model.safetensors') with open('/tmp/gpv2/adapter_config.json', 'w') as f: json.dump(cfg, f, indent=2) return kept # ---- GPU function ---- @spaces.GPU(duration=2400) def run_pr_v1(): log = [] def L(msg): log.append(msg) print(msg) t0 = time.time() L(f"[{time.strftime('%H:%M:%S')}] PR-v1 | GPU: {torch.cuda.get_device_name(0)}") L(f"VRAM: {torch.cuda.get_device_properties(0).total_memory/1e9:.1f} GB") # Build GPV-2 from mounted source kept = build_gpv2(f"{DATA_DIR}/gpv2") L(f"GPV-2 built: {kept} keys (18 = L30-32 MLP)") tok = AutoTokenizer.from_pretrained('Qwen/Qwen2.5-3B', trust_remote_code=True) if tok.pad_token is None: tok.pad_token = tok.eos_token def _eval(model, problems, n=30): code_ok = ans_ok = hyb_ok = total = 0 for p in problems[:n]: prompt = build_prompt(p) inp = tok(prompt, return_tensors='pt', truncation=True, max_length=1024) inp = {k: v.cuda() for k,v in inp.items()} with torch.no_grad(): out = model.generate(**inp, max_new_tokens=256, do_sample=False, pad_token_id=tok.eos_token_id) text = tok.decode(out[0][inp['input_ids'].shape[1]:], skip_special_tokens=True) m = re.search(r'(?i)FINAL_ANSWER\s*:\s*([^\n]+)', text) ans = m.group(1).strip().rstrip('.}"') if m else '' expected = str(p['expected_answer']).strip() total += 1 if ans == expected: ans_ok += 1 if 'def ' in text or 'import ' in text or '```' in text: code_ok += 1 if ans == expected: hyb_ok += 1 return (code_ok+ans_ok+hyb_ok)/(3*total), code_ok/total, ans_ok/total, hyb_ok/total # ---- Baseline ---- L("--- Frozen baseline ---") base = AutoModelForCausalLM.from_pretrained('Qwen/Qwen2.5-3B', torch_dtype=torch.float16, trust_remote_code=True).cuda() m_frozen = PeftModel.from_pretrained(base, '/tmp/gpv2').merge_and_unload() m_frozen.eval() comp_f, code_f, ans_f, hyb_f = _eval(m_frozen, all_probs, 30) L(f"FROZEN: comp={comp_f:.4f} code={code_f:.3f} ans={ans_f:.3f} hyb={hyb_f:.3f}") del m_frozen, base; gc.collect(); torch.cuda.empty_cache() L(f"VRAM free: {torch.cuda.memory_allocated()/1e9:.1f} GB") # ---- Training ---- L("--- Training 100 steps ---") repair_probs = all_probs[30:38] base_t = AutoModelForCausalLM.from_pretrained( 'Qwen/Qwen2.5-3B', torch_dtype=torch.float16, trust_remote_code=True, ).cuda() base_t.config.use_cache = False base_t.gradient_checkpointing_enable() model = PeftModel.from_pretrained(base_t, '/tmp/gpv2', is_trainable=True) model.enable_input_require_grads() model.train() trainable = sum(p.numel() for p in model.parameters() if p.requires_grad) L(f"Trainable: {trainable}") opt = AdamW([p for p in model.parameters() if p.requires_grad], lr=2e-5, foreach=False) MAX_LEN, N_STEPS = 96, 100 batches = [] for p in repair_probs: enc = tok(build_prompt(p), return_tensors='pt', truncation=True, max_length=MAX_LEN) batches.append(enc) L("Starting training...") losses = [] for step in range(N_STEPS): total_loss = 0.0 for enc in batches: inp = {k: v.cuda() for k,v in enc.items()} loss = model(**inp, labels=inp['input_ids']).loss total_loss += loss.item() loss.backward() opt.step() opt.zero_grad(set_to_none=True) losses.append(total_loss/len(batches)) if step % 10 == 0: torch.cuda.empty_cache() L(f" {step}: loss={losses[-1]:.4f}") L(f" final: loss={losses[-1]:.4f}") # ---- Repaired eval ---- L("--- Repaired eval ---") model = model.merge_and_unload() model.eval() comp_r, code_r, ans_r, hyb_r = _eval(model, all_probs, 30) delta = comp_r - comp_f L(f"REPAIRED: comp={comp_r:.4f} code={code_r:.3f} ans={ans_r:.3f} hyb={hyb_r:.3f}") L(f"DELTA: {delta:+.4f}") verdict = "PASS_REPAIR_SIGNAL" if delta >= 0.02 else "FAIL_LEARNED_REPAIR" L(f"VERDICT: {verdict}") elapsed = time.time() - t0 L(f"Done in {elapsed:.0f}s ({elapsed/60:.1f}m)") results = { "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"), "gpu": torch.cuda.get_device_name(0), "frozen": {"comp": comp_f, "code": code_f, "ans": ans_f, "hyb": hyb_f}, "repaired": {"comp": comp_r, "code": code_r, "ans": ans_r, "hyb": hyb_r}, "delta": delta, "verdict": verdict, "config": {"steps": N_STEPS, "lr": 2e-5, "max_len": MAX_LEN, "dtype": "float16"}, "elapsed_s": elapsed, } del model; gc.collect(); torch.cuda.empty_cache() return results, "\n".join(log) with gr.Blocks(title="PR-v1 Learned Repair") as demo: gr.Markdown("""# PR-v1: Learned Repair Pilot **Graft Planner v0.1 → Projection/Repair v1** 100 steps LoRA repair on GPV-2 (3B L30-32 MLP) GPU: RTX Pro 6000 (48GB) via ZeroGPU — ~30 min""") btn = gr.Button("🚀 Run PR-v1", variant="primary") json_out = gr.JSON(label="Results") log_out = gr.Textbox(label="Log", lines=22, max_lines=50) btn.click(fn=run_pr_v1, outputs=[json_out, log_out]) if __name__ == "__main__": demo.launch()