| """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" |
|
|
| |
| 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 |
|
|
| |
| @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") |
| |
| |
| 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 |
| |
| |
| 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") |
| |
| |
| 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}") |
| |
| |
| 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() |
|
|
