"""Train Qwen3-8B-Base with 40-pair recipe, eval on BigCodeBench-Hard. BigCodeBench is harder than HumanEval (real-world Python tasks, library use). Qwen3-8B-Base likely has headroom there (~30-45% baseline). Tests if recipe generalizes to newer model AND harder benchmark. """ import os, json, time, re, subprocess, tempfile, argparse os.environ.setdefault("HF_HOME", "/workspace/hf") os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1") os.environ["TRANSFORMERS_VERBOSITY"] = "error" import torch from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer from datasets import load_dataset, Dataset as HFDataset from peft import LoraConfig, get_peft_model T0 = time.time() def log(m): print(f"[{time.time()-T0:7.1f}s] {m}", flush=True) def extract_code(text): if "```python" in text: text = text.split("```python", 1)[1] elif "```" in text: text = text.split("```", 1)[1] if "```" in text: text = text.split("```", 1)[0] return text.strip() def verify_bcb(code, test_code): runner = "\n\nif __name__ == '__main__':\n import unittest; unittest.main(argv=['x'], exit=False, verbosity=0)\n" body = code + "\n\n" + test_code + runner with tempfile.NamedTemporaryFile("w", suffix=".py", delete=False) as f: f.write(body); path = f.name try: r = subprocess.run(["python3", path], capture_output=True, timeout=20, text=True, cwd="/tmp") out = (r.stdout or "") + "\n" + (r.stderr or "") if "OK" in out and "FAILED" not in out and "Error" not in out and r.returncode == 0: return True return False except subprocess.TimeoutExpired: return False finally: try: os.unlink(path) except: pass def gen_batch(model, tok, prompts, max_new=600, temperature=0.0, batch=4): outs = [] for i in range(0, len(prompts), batch): chunk = prompts[i:i+batch] texts = [] for p in chunk: msgs = [{"role": "system", "content": "You are an expert Python coder. Output one ```python block with the complete solution."}, {"role": "user", "content": p}] texts.append(tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)) inp = tok(texts, return_tensors="pt", padding=True, truncation=True, max_length=2000).to(model.device) with torch.no_grad(): out = model.generate(**inp, max_new_tokens=max_new, do_sample=temperature > 0, temperature=temperature if temperature > 0 else 1.0, top_p=0.95, pad_token_id=tok.eos_token_id) for j in range(out.size(0)): outs.append(tok.decode(out[j][inp.input_ids.shape[1]:], skip_special_tokens=True)) return outs def eval_bcb_hard(model, tok, label, max_n=148): bcb = list(load_dataset("bigcode/bigcodebench-hard", split="v0.1.4"))[:max_n] log(f" BCB-Hard [{label}] ({len(bcb)})") prompts = [p["instruct_prompt"] for p in bcb] outs = gen_batch(model, tok, prompts, max_new=700, batch=4) correct = 0 for i, (p, raw) in enumerate(zip(bcb, outs)): code = extract_code(raw) if "```" in raw else raw if verify_bcb(code, p["test"]): correct += 1 if (i+1) % 20 == 0: log(f" {label} BCB {i+1}/{len(bcb)}: {correct}") return correct, len(bcb) def eval_humaneval(model, tok, label): he = list(load_dataset("openai_humaneval", split="test")) log(f" HumanEval [{label}] ({len(he)})") prompts = [p["prompt"] + "\n# Complete the function above." for p in he] outs = gen_batch(model, tok, prompts, max_new=400, batch=4) correct = 0 for i, (p, raw) in enumerate(zip(he, outs)): code = extract_code(raw) if "```" in raw else raw full = p["prompt"] + "\n" + code if "def " not in code else code test_code = full + "\n\n" + p["test"] + f"\n\ncheck({p['entry_point']})" with tempfile.NamedTemporaryFile("w", suffix=".py", delete=False) as f: f.write(test_code); path = f.name try: r = subprocess.run(["python3", path], capture_output=True, timeout=10, text=True, cwd="/tmp") if r.returncode == 0: correct += 1 except subprocess.TimeoutExpired: pass finally: try: os.unlink(path) except: pass if (i+1) % 40 == 0: log(f" {label} HE {i+1}/{len(he)}: {correct}") return correct, len(he) def make_example(r, tok): user = (f"Implement: {r['signature']}\n\n" f"Tests:\n{chr(10).join(r['tests'])}\n\n" f"My attempt:\n```python\n{r['broken']}\n```\n\n" f"Error:\n{r.get('error','')}\n\n" f"Fix and output the corrected code only.") assistant = f"```python\n{r['fixed']}\n```" msgs_pre = [{"role": "system", "content": "You are an expert Python coder. Output one ```python block with the complete solution."}, {"role": "user", "content": user}] msgs_full = msgs_pre + [{"role": "assistant", "content": assistant}] pre = tok.apply_chat_template(msgs_pre, tokenize=False, add_generation_prompt=True) full = tok.apply_chat_template(msgs_full, tokenize=False) pre_ids = tok(pre, add_special_tokens=False)["input_ids"] full_ids = tok(full, add_special_tokens=False)["input_ids"] MAX = 1024 full_ids = full_ids[:MAX] labels = list(full_ids) n_pre = min(len(pre_ids), len(labels)) for i in range(n_pre): labels[i] = -100 pad = MAX - len(full_ids) return {"input_ids": full_ids + [tok.pad_token_id]*pad, "attention_mask": [1]*len(full_ids) + [0]*pad, "labels": labels + [-100]*pad} def main(): ap = argparse.ArgumentParser() ap.add_argument("--model", required=True) ap.add_argument("--pairs", default="/workspace/saved_pairs/pairs_40.jsonl") ap.add_argument("--n_pairs", type=int, default=40) ap.add_argument("--tag", required=True) args = ap.parse_args() out_dir = f"/workspace/bcb_eval/{args.tag}" os.makedirs(out_dir, exist_ok=True) log(f"loading {args.model}") tok = AutoTokenizer.from_pretrained(args.model) if tok.pad_token is None: tok.pad_token = tok.eos_token tok.padding_side = "left" model = AutoModelForCausalLM.from_pretrained(args.model, torch_dtype=torch.bfloat16, device_map="cuda:0") log(f" loaded mem={torch.cuda.memory_allocated('cuda:0')/1e9:.1f}GB") model.eval() log("=== BASE evals ===") base_he, _ = eval_humaneval(model, tok, "BASE") base_bcb, _ = eval_bcb_hard(model, tok, "BASE") log(f" BASE: HumanEval={base_he}/164 BCB-Hard={base_bcb}/148") pairs = [json.loads(l) for l in open(args.pairs)][:args.n_pairs] log(f"=== TRAINING — {len(pairs)} pairs ===") lora_cfg = LoraConfig(r=16, lora_alpha=32, lora_dropout=0.05, bias="none", target_modules=["q_proj", "k_proj", "v_proj", "o_proj"], task_type="CAUSAL_LM") model = get_peft_model(model, lora_cfg) tok.padding_side = "right" ds = HFDataset.from_list([make_example(r, tok) for r in pairs]) targs = TrainingArguments( output_dir=f"{out_dir}/ckpt", num_train_epochs=2, per_device_train_batch_size=1, gradient_accumulation_steps=4, learning_rate=1e-4, bf16=True, logging_steps=10, save_strategy="no", report_to="none", remove_unused_columns=False, warmup_ratio=0.05, ) Trainer(model=model, args=targs, train_dataset=ds, processing_class=tok).train() log(" training done") tok.padding_side = "left" model.eval() log("=== TRAINED evals ===") tr_he, _ = eval_humaneval(model, tok, "TRAINED") tr_bcb, _ = eval_bcb_hard(model, tok, "TRAINED") result = { "model": args.model, "method": "warmup 40 pairs", "humaneval": {"base": base_he, "trained": tr_he, "delta": tr_he-base_he, "n": 164}, "bcb_hard": {"base": base_bcb, "trained": tr_bcb, "delta": tr_bcb-base_bcb, "n": 148}, "elapsed_s": time.time() - T0, } with open(f"{out_dir}/result.json", "w") as fh: json.dump(result, fh, indent=2) print() print("=" * 70) print(f" {args.model}") print(f" HumanEval: base={base_he}/164 trained={tr_he}/164 Δ={tr_he-base_he:+d}") print(f" BCB-Hard: base={base_bcb}/148 trained={tr_bcb}/148 Δ={tr_bcb-base_bcb:+d}") print(f" Time: {time.time()-T0:.0f}s") print("=" * 70) if __name__ == "__main__": main()