tinyforge-zero/recipe/bootstrap.py

"""Self-Bootstrapping TinyForge.

Single model. No external dataset. Just a Python interpreter.

Loop:
  for iter in 1..N:
    1. Model generates K problems (function signature + tests + canonical solution)
    2. Filter: keep only those where canonical executes & tests pass
    3. Model solves each fresh (forget canonical)
    4. Verify against tests → identify failures
    5. Model repairs each failure (one shot, with error)
    6. Verify repairs → collect (broken, fixed) pairs
    7. Periodically: LoRA-train on accumulated pairs
    8. Periodically: eval on held-out HumanEval-mini

If accuracy on HumanEval rises without ever seeing HumanEval problems → recipe works.
"""
import os, sys, json, time, re, gc, subprocess, tempfile, argparse, random, math
os.environ.setdefault("HF_HOME", "/workspace/hf")
os.environ["TRANSFORMERS_VERBOSITY"] = "error"
os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"

import torch
from torch.utils.data import Dataset, DataLoader
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 run_python(code, timeout=8):
    """Run code in subprocess. Return (passed, stderr_or_msg)."""
    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")
        if r.returncode == 0: return True, ""
        err = (r.stderr or r.stdout).strip().splitlines()
        return False, "\n".join(err[-3:])[:300]
    except subprocess.TimeoutExpired: return False, "timeout"
    finally:
        try: os.unlink(path)
        except: pass


def gen_batch(model, tok, prompts, max_new=400, temperature=0.7, batch=8):
    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 a Python coder."},
                    {"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=1500).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


PROBLEM_GEN_PROMPT = """Generate ONE simple Python coding problem with a clear function spec and 3 test assertions.

Output format (exactly one ```python block):

```python
def {function_name}({args}):
    \"\"\"{one-line description of what the function does}\"\"\"
    {implementation}

# tests
assert {function_name}(...) == ...
assert {function_name}(...) == ...
assert {function_name}(...) == ...
```

Make the function specific and concrete. The function should be 3-15 lines. Tests must verify the function works correctly. Output ONLY the code block."""


def parse_generated_problem(raw_code):
    """Split into (function_signature_with_docstring, full_solution_code, test_lines).
    Returns None if parsing fails or it's malformed."""
    code = raw_code.strip()
    if "def " not in code: return None

    # Find first def
    lines = code.split("\n")
    func_start = None
    for i, l in enumerate(lines):
        if l.startswith("def "):
            func_start = i; break
    if func_start is None: return None

    # Find tests (assert lines after the def block)
    tests = []
    in_def_body = False
    def_end = None
    for i in range(func_start, len(lines)):
        l = lines[i]
        if l.startswith("def ") and i > func_start: break
        if l.startswith("assert "):
            tests.append(l)
            if def_end is None: def_end = i
        elif tests and not l.strip().startswith(("#", "assert", "")):
            break

    if len(tests) < 2: return None
    if def_end is None: def_end = len(lines)

    full_solution = "\n".join(lines[func_start:def_end]).strip()
    if len(full_solution) < 30: return None

    # Build function signature stub for re-implementation
    # Find docstring if present
    sig_lines = []
    for i in range(func_start, def_end):
        l = lines[i]
        sig_lines.append(l)
        if i > func_start and l.strip().endswith('"""') and ('"""' in lines[i-1] or '"""' in l[:l.rfind('"""')]):
            break
        if i > func_start and l.strip().startswith('"""') and l.strip().endswith('"""') and l.strip() != '"""':
            break
        # If no docstring, stop after the def line itself
        if i == func_start and not any('"""' in lines[j] for j in range(i, min(i+5, def_end))):
            sig_lines.append("    pass")
            break
    signature = "\n".join(sig_lines)

    # Extract function name from signature
    m = re.match(r"def\s+(\w+)\s*\(", lines[func_start])
    if not m: return None
    fn_name = m.group(1)

    return {
        "fn_name": fn_name,
        "signature": signature,
        "canonical": full_solution,
        "tests": tests,
        "raw": code,
    }


# ── Loop ────────────────────────────────────────────────────────────────

def humaneval_eval(model, tok, n=30):
    """Eval on HumanEval-mini (first N problems)."""
    he = list(load_dataset("openai_humaneval", split="test"))[:n]
    prompts = [p["prompt"] + "\n# Complete the function above." for p in he]
    outs = gen_batch(model, tok, prompts, max_new=400, temperature=0.0, batch=4)
    correct = 0
    for p, raw in zip(he, outs):
        code = extract_code(raw) if "```" in raw else raw
        # Try the model's completion combined with the prompt
        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']})"
        ok, _ = run_python(test_code, timeout=10)
        if ok: correct += 1
    return correct, n


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", default="Qwen/Qwen2.5-Coder-1.5B-Instruct")
    ap.add_argument("--gpu", type=int, default=0)
    ap.add_argument("--iterations", type=int, default=20)
    ap.add_argument("--problems_per_iter", type=int, default=16)
    ap.add_argument("--train_every", type=int, default=10)
    ap.add_argument("--eval_every", type=int, default=10)
    ap.add_argument("--tag", required=True)
    args = ap.parse_args()

    out_dir = f"/workspace/bootstrap/{args.tag}"
    os.makedirs(out_dir, exist_ok=True)
    device = torch.device(f"cuda:{args.gpu}")

    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, dtype=torch.bfloat16, device_map=f"cuda:{args.gpu}")
    log(f"  loaded mem={torch.cuda.memory_allocated(device)/1e9:.1f}GB")

    # Initial eval
    log("INITIAL eval on HumanEval-mini")
    init_correct, init_total = humaneval_eval(model, tok, n=30)
    log(f"  HumanEval-mini base: {init_correct}/{init_total}")

    # LoRA setup (will be applied for training, base kept frozen)
    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)
    log(f"  LoRA applied; trainable={sum(p.numel() for p in model.parameters() if p.requires_grad)/1e6:.1f}M")

    accumulated_pairs = []
    eval_log = [{"iter": 0, "correct": init_correct, "total": init_total}]
    iter_stats = []

    for it in range(1, args.iterations + 1):
        it_t = time.time()
        # 1. Generate K problems
        gen_prompts = [PROBLEM_GEN_PROMPT for _ in range(args.problems_per_iter)]
        raw_problems = gen_batch(model, tok, gen_prompts, max_new=400, temperature=0.9)

        # 2. Parse + verify canonical
        valid_problems = []
        for raw in raw_problems:
            code = extract_code(raw) if "```" in raw else raw
            parsed = parse_generated_problem(code)
            if parsed is None: continue
            # Verify canonical passes its own tests
            full = parsed["canonical"] + "\n\n" + "\n".join(parsed["tests"])
            ok, _ = run_python(full)
            if ok: valid_problems.append(parsed)

        if not valid_problems:
            log(f"iter {it}: 0 valid problems generated, skipping")
            iter_stats.append({"iter": it, "valid": 0, "fails": 0, "repairs": 0})
            continue

        # 3. Model solves each fresh — N=4 sampled attempts at temp=0.8 to surface natural fails
        N_ATTEMPTS = 4
        solve_prompts = [f"Implement this function so it passes the tests below.\n\n```python\n{p['signature']}\n```\n\nTests:\n{chr(10).join(p['tests'])}\n\nOutput only the function implementation in one ```python block." for p in valid_problems]
        # Generate N attempts each (4 * len(prompts) total)
        all_solve_prompts = solve_prompts * N_ATTEMPTS
        all_attempts = gen_batch(model, tok, all_solve_prompts, max_new=400, temperature=0.8)
        # Reshape: by problem, list of N attempts
        per_problem_attempts = [all_attempts[i::len(valid_problems)] for i in range(len(valid_problems))]

        # 4-5. Mine (broken, fixed) pairs from same model's diverse outputs
        failures = []
        new_pairs = 0
        for p, attempts in zip(valid_problems, per_problem_attempts):
            broken_one = None; fixed_one = None; broken_err = None
            for raw in attempts:
                code = extract_code(raw) if "```" in raw else raw
                full = code + "\n\n" + "\n".join(p["tests"])
                ok, err = run_python(full)
                if ok and fixed_one is None:
                    fixed_one = code
                elif not ok and broken_one is None:
                    broken_one = code; broken_err = err
                if broken_one and fixed_one: break
            if broken_one is None:
                continue
            if fixed_one is not None:
                # Self-mined repair pair from same-model diverse outputs
                accumulated_pairs.append({
                    "signature": p["signature"], "tests": p["tests"],
                    "broken": broken_one, "error": broken_err, "fixed": fixed_one,
                })
                new_pairs += 1
            else:
                # All attempts failed — try one more repair pass with explicit error
                failures.append({"p": p, "broken": broken_one, "error": broken_err})

        # Optional: try repair on remaining all-failed cases
        if failures:
            repair_prompts = [f"Implement: {f['p']['signature']}\n\nTests:\n{chr(10).join(f['p']['tests'])}\n\nMy attempt:\n```python\n{f['broken']}\n```\n\nError:\n{f['error']}\n\nFix and output the corrected code only." for f in failures]
            repairs = gen_batch(model, tok, repair_prompts, max_new=400, temperature=0.8)
            for f, raw in zip(failures, repairs):
                fix = extract_code(raw) if "```" in raw else raw
                full = fix + "\n\n" + "\n".join(f["p"]["tests"])
                ok, _ = run_python(full)
                if ok:
                    accumulated_pairs.append({
                        "signature": f["p"]["signature"], "tests": f["p"]["tests"],
                        "broken": f["broken"], "error": f["error"], "fixed": fix,
                    })
                    new_pairs += 1

        log(f"iter {it}: {len(valid_problems)} valid problems, {len(failures)} failures, {new_pairs} repair pairs harvested (total: {len(accumulated_pairs)})  [{time.time()-it_t:.0f}s]")
        iter_stats.append({"iter": it, "valid": len(valid_problems), "fails": len(failures), "repairs": new_pairs, "elapsed": time.time()-it_t})

        # Save incrementally (in case of crash)
        with open(f"{out_dir}/pairs.jsonl", "w") as fh:
            for r in accumulated_pairs: fh.write(json.dumps(r) + "\n")

        # 6. Periodic training
        if it % args.train_every == 0 and len(accumulated_pairs) >= 10:
            log(f"  TRAINING on {len(accumulated_pairs)} pairs")
            tok.padding_side = "right"

            def make_example(r):
                user = f"Implement: {r['signature']}\n\nTests:\n{chr(10).join(r['tests'])}\n\nMy attempt:\n```python\n{r['broken']}\n```\n\nError:\n{r['error']}\n\nFix and output the corrected code only."
                assistant = f"```python\n{r['fixed']}\n```"
                msgs_pre = [{"role": "system", "content": "You are a Python coder."},
                            {"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}

            ds = HFDataset.from_list([make_example(r) for r in accumulated_pairs])
            targs = TrainingArguments(
                output_dir=f"{out_dir}/ckpt_iter{it}", num_train_epochs=2,
                per_device_train_batch_size=1, gradient_accumulation_steps=4,
                learning_rate=1e-4, bf16=True, logging_steps=20,
                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()
            tok.padding_side = "left"

        # 7. Periodic eval
        if it % args.eval_every == 0:
            model.eval()
            corr, tot = humaneval_eval(model, tok, n=30)
            log(f"  HumanEval-mini @ iter {it}: {corr}/{tot}")
            eval_log.append({"iter": it, "correct": corr, "total": tot})
            model.train()

    # Final eval
    model.eval()
    final_correct, final_total = humaneval_eval(model, tok, n=30)
    eval_log.append({"iter": args.iterations, "correct": final_correct, "total": final_total, "final": True})

    # Save everything
    with open(f"{out_dir}/iter_stats.jsonl", "w") as fh:
        for r in iter_stats: fh.write(json.dumps(r) + "\n")
    with open(f"{out_dir}/eval_log.json", "w") as fh:
        json.dump(eval_log, fh, indent=2)
    with open(f"{out_dir}/pairs.jsonl", "w") as fh:
        for r in accumulated_pairs: fh.write(json.dumps(r) + "\n")

    print()
    print("=" * 70)
    print(f"  MODEL: {args.model}")
    print(f"  ITERATIONS: {args.iterations}, problems/iter: {args.problems_per_iter}")
    print(f"  TOTAL repair pairs: {len(accumulated_pairs)}")
    print(f"  HUMANEVAL-MINI:  base={init_correct}/{init_total}  final={final_correct}/{final_total}  Δ={final_correct-init_correct:+d}")
    print(f"  time: {time.time()-T0:.0f}s")
    print("=" * 70)


if __name__ == "__main__":
    main()