tinyforge-zero/evals/eval_raw.py

"""vLLM dual eval using RAW completion format (no chat template) for base models.

Recipe for non-instruct base models — uses simple completion-style prompting
that matches how base models were pretrained.
"""
import os, json, time, re, subprocess, tempfile, argparse, gc
os.environ.setdefault("HF_HOME", "/workspace/hf")
os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")
os.environ["TRANSFORMERS_VERBOSITY"] = "error"

import torch
from datasets import load_dataset

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 run_python(code, timeout=10):
    with tempfile.NamedTemporaryFile("w", suffix=".py", delete=False) as f:
        f.write(code); path = f.name
    try:
        r = subprocess.run(["python3", path], capture_output=True, timeout=timeout, text=True, cwd="/tmp")
        return r.returncode == 0
    except subprocess.TimeoutExpired: return False
    finally:
        try: os.unlink(path)
        except: pass


def make_he_prompt(p):
    """Raw completion: just the docstring + 'def'."""
    return p["prompt"]


def make_mbpp_prompt(p):
    """Raw completion: docstring + tests + 'def'."""
    return (f"# Task: {p['prompt']}\n"
            f"# Tests:\n# " + "\n# ".join(p["test_list"]) + "\n\n")


def vllm_generate(llm, prompts, max_new=400, temperature=0.0, stops=None):
    from vllm import SamplingParams
    sp = SamplingParams(
        temperature=temperature, top_p=0.95 if temperature > 0 else 1.0,
        max_tokens=max_new, stop=stops or ["\nclass ", "\nif __name__", "\nprint(", "\n#"],
    )
    out = llm.generate(prompts, sp, use_tqdm=False)
    return [o.outputs[0].text for o in out]


def vllm_generate_lora(llm, prompts, lora_req, max_new=400, temperature=0.0, stops=None):
    from vllm import SamplingParams
    sp = SamplingParams(
        temperature=temperature, top_p=0.95 if temperature > 0 else 1.0,
        max_tokens=max_new, stop=stops or ["\nclass ", "\nif __name__", "\nprint(", "\n#"],
    )
    out = llm.generate(prompts, sp, lora_request=lora_req, use_tqdm=False)
    return [o.outputs[0].text for o in out]


def eval_humaneval(outs_func, label):
    he = list(load_dataset("openai_humaneval", split="test"))
    log(f"  HumanEval [{label}] ({len(he)})")
    prompts = [make_he_prompt(p) for p in he]
    t0 = time.time()
    outs = outs_func(prompts, max_new=400)
    log(f"    gen done in {time.time()-t0:.1f}s")
    correct = 0
    for p, raw in zip(he, outs):
        # construct full function: prompt + raw completion
        full = p["prompt"] + raw
        test_code = full + "\n\n" + p["test"] + f"\n\ncheck({p['entry_point']})"
        if run_python(test_code, timeout=10): correct += 1
    return correct, len(he)


def eval_mbpp(outs_func, label, n=200):
    mbpp = list(load_dataset("mbpp", "sanitized", split="test"))[:n]
    log(f"  MBPP [{label}] ({len(mbpp)})")
    prompts = [make_mbpp_prompt(p) for p in mbpp]
    t0 = time.time()
    outs = outs_func(prompts, max_new=400)
    log(f"    gen done in {time.time()-t0:.1f}s")
    correct = 0
    for p, raw in zip(mbpp, outs):
        # raw is the function code
        code = raw
        if "```" in code:
            code = extract_code("```python" + code if "```python" not in code else code)
        test_code = code + "\n\n" + "\n".join(p["test_list"])
        if run_python(test_code, timeout=10): correct += 1
    return correct, len(mbpp)


def make_train_example(r, tok):
    """Raw-completion training format."""
    sig = r.get("signature", "")
    broken = r.get("broken", "")
    fixed = r.get("fixed", "")
    tests = r.get("tests", [])
    err = r.get("error", "")
    user = (f"# Task: implement {sig}\n"
            f"# Tests:\n# " + "\n# ".join(tests) + "\n"
            f"# My broken attempt:\n{broken}\n"
            f"# Error: {err}\n"
            f"# Corrected:\n")
    target = fixed
    full = user + target
    full_ids = tok(full, add_special_tokens=False)["input_ids"]
    user_ids = tok(user, add_special_tokens=False)["input_ids"]
    MAX = 1024
    full_ids = full_ids[:MAX]
    labels = list(full_ids)
    n_user = min(len(user_ids), len(labels))
    for i in range(n_user): 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("--mbpp_n", type=int, default=200)
    ap.add_argument("--tag", required=True)
    ap.add_argument("--skip_train", action="store_true")
    args = ap.parse_args()

    out_dir = f"/workspace/dual_eval_raw/{args.tag}"
    os.makedirs(out_dir, exist_ok=True)

    from vllm import LLM
    from transformers import AutoTokenizer
    log(f"loading {args.model} into vLLM")
    tok = AutoTokenizer.from_pretrained(args.model)
    if tok.pad_token is None: tok.pad_token = tok.eos_token
    llm = LLM(model=args.model, dtype="bfloat16", gpu_memory_utilization=0.85, max_model_len=2048)
    log(f"  loaded")

    log("=== BASE evals ===")
    base_he, _ = eval_humaneval(lambda P, max_new=400: vllm_generate(llm, P, max_new=max_new), "BASE")
    base_mbpp, _ = eval_mbpp(lambda P, max_new=400: vllm_generate(llm, P, max_new=max_new), "BASE", n=args.mbpp_n)
    log(f"  BASE: HumanEval={base_he}/164  MBPP={base_mbpp}/{args.mbpp_n}")

    if args.skip_train:
        result = {"model": args.model, "base_humaneval": base_he, "base_mbpp": base_mbpp, "n_he": 164, "n_mbpp": args.mbpp_n, "elapsed_s": time.time()-T0}
        with open(f"{out_dir}/result.json", "w") as fh: json.dump(result, fh, indent=2)
        return

    # Tear down vLLM, train LoRA
    log("=== TRAINING ===")
    del llm; gc.collect(); torch.cuda.empty_cache()

    from transformers import AutoModelForCausalLM, TrainingArguments, Trainer
    from datasets import Dataset as HFDataset
    from peft import LoraConfig, get_peft_model

    pairs = [json.loads(l) for l in open(args.pairs)][:args.n_pairs]
    model = AutoModelForCausalLM.from_pretrained(args.model, torch_dtype=torch.bfloat16, device_map="cuda:0")
    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)

    ds = HFDataset.from_list([make_train_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, tokenizer=tok).train()
    log("training done")

    adapter_dir = f"{out_dir}/adapter"
    model.save_pretrained(adapter_dir)
    del model; gc.collect(); torch.cuda.empty_cache()

    from vllm import LLM
    from vllm.lora.request import LoRARequest
    llm = LLM(model=args.model, dtype="bfloat16", gpu_memory_utilization=0.85, max_model_len=2048,
              enable_lora=True, max_lora_rank=16)
    lora_req = LoRARequest("tf_adapter", 1, adapter_dir)

    log("=== TRAINED evals (vLLM + LoRA) ===")
    tr_he, _ = eval_humaneval(lambda P, max_new=400: vllm_generate_lora(llm, P, lora_req, max_new=max_new), "TRAINED")
    tr_mbpp, _ = eval_mbpp(lambda P, max_new=400: vllm_generate_lora(llm, P, lora_req, max_new=max_new), "TRAINED", n=args.mbpp_n)

    result = {
        "model": args.model, "n_pairs": len(pairs),
        "humaneval": {"base": base_he, "trained": tr_he, "delta": tr_he-base_he, "n": 164},
        "mbpp": {"base": base_mbpp, "trained": tr_mbpp, "delta": tr_mbpp-base_mbpp, "n": args.mbpp_n},
        "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} — RAW completion format")
    print(f"  HumanEval:  base={base_he}/164  trained={tr_he}/164  Δ={tr_he-base_he:+d}")
    print(f"  MBPP:       base={base_mbpp}/{args.mbpp_n}  trained={tr_mbpp}/{args.mbpp_n}  Δ={tr_mbpp-base_mbpp:+d}")
    print(f"  Time: {time.time()-T0:.0f}s")
    print("=" * 70)


if __name__ == "__main__":
    main()