apunkt/tinyforge-zero
1
0
Fork 0
mirror of https://github.com/ranausmanai/tinyforge-zero.git synced 2026-06-11 21:05:12 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions

Ship every paper-referenced experiment script

Browse source 
Reorganizes the repo so every section of the paper has a corresponding
script. Previously only the core recipe + control + evals were here.

New subdirs:
- tts/             — test-time sampling (§2.2, §3.3): scaling sweep, HE, MATH-500,
                     AIME, 14B-recipe + TTS, 8B-raw-TTS control.
- experiments/     — every §3 finding as a runnable script:
                     · self_consistency (§3.4)
                     · recipe_x_tts_synergy (§3.5, novel)
                     · mbpp_seeded_cross_arch (§3.9)
                     · cross_domain_code_to_math (§3.10)
                     · self_correction_math_{naive,fixed} (§3.10, the
                       catastrophic-then-recovered case)
                     · math500_seeded_mining (§3.10 distribution mismatch)
                     · bcb_hard_eval (§3.10 distribution mismatch)
                     · recursive_bootstrap (§3.10 plateau)
                     · diversity_cued_mining (§3.10 low yield)
                     · aime_scaling (TTS curve)
                     · star_baseline_gsm8k (related-work baseline)
- evals/           — moved out of recipe/ (eval_raw, eval_plus, confirm)

Also adds: bootstrap_14b_4bit_harvest, curriculum_code, math_bootstrap to
recipe/ for completeness.

REPRODUCE.md now maps each paper section / table / figure to its exact
script and expected output.
This commit is contained in:
Rana Usman 2026-05-13 21:09:54 +05:00
parent c867697f7c
commit 826f934d2e
27 changed files with 4467 additions and 134 deletions
Download patch file Download diff file Expand all files Collapse all files

191
recipe/bootstrap_14b_4bit_harvest.py Normal file
View file

@ -0,0 +1,191 @@
"""Bootstrap loop adapted for large models — uses 4-bit NF4 quantization and batch=1.
Just the harvest loop (no training during loop). Saves pairs.
"""
import os, sys, json, time, re, gc, subprocess, tempfile, argparse, random
os.environ.setdefault("HF_HOME", "/workspace/hf")
os.environ.setdefault("CUDA_VISIBLE_DEVICES", "1")
os.environ["TRANSFORMERS_VERBOSITY"] = "error"
os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
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=8):
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_one(model, tok, prompt, max_new=400, temperature=0.0):
msgs = [{"role": "system", "content": "You are a Python coder."},
{"role": "user", "content": prompt}]
text = tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
inp = tok(text, return_tensors="pt", 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)
return tok.decode(out[0][inp.input_ids.shape[1]:], skip_special_tokens=True)
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. Output ONLY the code block."""
def parse_problem(raw_code):
code = raw_code.strip()
if "def " not in code: return None
lines = code.split("\n")
func_start = next((i for i, l in enumerate(lines) if l.startswith("def ")), None)
if func_start is None: return None
tests = []
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
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
m = re.match(r"def\s+(\w+)\s*\(", lines[func_start])
if not m: return None
sig_lines = []
for i in range(func_start, def_end):
sig_lines.append(lines[i])
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
if i > func_start and '"""' in lines[i] and ('"""' in lines[i-1] or lines[i].count('"""') >= 2):
break
return {"fn_name": m.group(1), "signature": "\n".join(sig_lines), "tests": tests, "canonical": full_solution}
def humaneval_full(model, tok):
he = list(load_dataset("openai_humaneval", split="test"))
log(f" full HumanEval: {len(he)} problems")
correct = 0
for i, p in enumerate(he):
prompt = p["prompt"] + "\n# Complete the function above."
raw = gen_one(model, tok, prompt, max_new=400, temperature=0.0)
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']})"
ok, _ = run_python(test_code, timeout=10)
if ok: correct += 1
if (i+1) % 20 == 0: log(f" eval {i+1}/{len(he)}: {correct} correct")
return correct, len(he)
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--model", default="Qwen/Qwen2.5-14B")
ap.add_argument("--iterations", type=int, default=20)
ap.add_argument("--problems_per_iter", type=int, default=8)
ap.add_argument("--n_attempts", type=int, default=4)
ap.add_argument("--tag", required=True)
args = ap.parse_args()
out_dir = f"/workspace/bootstrap14b/{args.tag}"
os.makedirs(out_dir, exist_ok=True)
log(f"loading {args.model} in 4-bit NF4")
tok = AutoTokenizer.from_pretrained(args.model)
if tok.pad_token is None: tok.pad_token = tok.eos_token
tok.padding_side = "left"
bnb_cfg = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16,
bnb_4bit_use_double_quant=True)
model = AutoModelForCausalLM.from_pretrained(args.model, quantization_config=bnb_cfg,
device_map="cuda:0")
model.eval()
log(f" loaded mem={torch.cuda.memory_allocated('cuda:0')/1e9:.1f}GB")
log("INITIAL eval on full HumanEval")
base_correct, base_total = humaneval_full(model, tok)
log(f" base: {base_correct}/{base_total}")
accumulated = []
for it in range(1, args.iterations + 1):
it_t = time.time()
valid_problems = []
for _ in range(args.problems_per_iter):
raw = gen_one(model, tok, PROBLEM_GEN_PROMPT, max_new=400, temperature=0.9)
code = extract_code(raw) if "```" in raw else raw
parsed = parse_problem(code)
if not parsed: continue
full = parsed["canonical"] + "\n\n" + "\n".join(parsed["tests"])
ok, _ = run_python(full)
if ok: valid_problems.append(parsed)
new_pairs = 0
for p in valid_problems:
attempts = []
solve_prompt = f"Implement: {p['signature']}\n\nTests:\n{chr(10).join(p['tests'])}\n\nOutput only the function implementation in one ```python block."
for _ in range(args.n_attempts):
raw = gen_one(model, tok, solve_prompt, max_new=400, temperature=0.8)
attempts.append(raw)
broken = None; fixed = 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 is None: fixed = code
elif not ok and broken is None: broken = code; broken_err = err
if broken and fixed: break
if broken and fixed:
accumulated.append({"signature": p["signature"], "tests": p["tests"],
"broken": broken, "error": broken_err if 'broken_err' in dir() else "",
"fixed": fixed})
new_pairs += 1
log(f"iter {it}: {len(valid_problems)} valid, {new_pairs} pairs (total: {len(accumulated)}) [{time.time()-it_t:.0f}s]")
with open(f"{out_dir}/pairs.jsonl", "w") as fh:
for r in accumulated: fh.write(json.dumps(r) + "\n")
log(f"DONE — accumulated {len(accumulated)} pairs from {args.iterations} iters")
print()
print("=" * 70)
print(f" 14B BASELINE: {base_correct}/{base_total} on HumanEval")
print(f" Accumulated pairs: {len(accumulated)}")
print(f" Time: {time.time()-T0:.0f}s")
print("=" * 70)
if __name__ == "__main__":
main()

165
recipe/confirm.py
View file

@ -1,165 +0,0 @@
"""Confirm the peak +5 result on full HumanEval (164 problems) and try the cliff at 39 pairs."""
import os, sys, json, time, re, gc, subprocess, tempfile, argparse
os.environ.setdefault("HF_HOME", "/workspace/hf")
os.environ.setdefault("CUDA_VISIBLE_DEVICES", "0")
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 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 gen_batch(model, tok, prompts, max_new=400, 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 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
def humaneval_full(model, tok):
he = list(load_dataset("openai_humaneval", split="test"))
log(f" full HumanEval: {len(he)} problems")
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
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']})"
if run_python(test_code, timeout=10): correct += 1
return correct, len(he)
def make_example(r, tok):
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}
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--n_pairs", type=int, default=21, help="how many pairs from the saved set to train on")
ap.add_argument("--epochs", type=int, default=2)
ap.add_argument("--seed", type=int, default=42)
ap.add_argument("--tag", required=True)
args = ap.parse_args()
torch.manual_seed(args.seed)
pairs_path = "/workspace/bootstrap/bs_7b_v3/pairs.jsonl"
pairs = [json.loads(l) for l in open(pairs_path)]
log(f"loaded {len(pairs)} pairs from prior bootstrap run")
pairs_use = pairs[:args.n_pairs]
log(f"using {len(pairs_use)} for this run")
out_dir = f"/workspace/confirm/{args.tag}"
os.makedirs(out_dir, exist_ok=True)
log("loading Qwen/Qwen2.5-7B")
tok = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B")
if tok.pad_token is None: tok.pad_token = tok.eos_token
tok.padding_side = "left"
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-7B", dtype=torch.bfloat16, device_map="cuda:0")
# Eval base
model.eval()
log("eval BASE on full HumanEval")
base_corr, base_total = humaneval_full(model, tok)
log(f" BASE: {base_corr}/{base_total}")
# Apply LoRA + train
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("LoRA applied")
tok.padding_side = "right"
examples = [make_example(r, tok) for r in pairs_use]
ds = HFDataset.from_list(examples)
targs = TrainingArguments(
output_dir=f"{out_dir}/ckpt", num_train_epochs=args.epochs,
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,
seed=args.seed,
)
log(f"training on {len(ds)} pairs, {args.epochs} epochs")
Trainer(model=model, args=targs, train_dataset=ds, processing_class=tok).train()
log("training done")
tok.padding_side = "left"
# Eval trained
model.eval()
log("eval TRAINED on full HumanEval")
tr_corr, tr_total = humaneval_full(model, tok)
log(f" TRAINED: {tr_corr}/{tr_total}")
result = {
"n_pairs_used": len(pairs_use), "epochs": args.epochs, "seed": args.seed,
"base": [base_corr, base_total], "trained": [tr_corr, tr_total],
"delta": tr_corr - base_corr,
"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" N_PAIRS: {len(pairs_use)} EPOCHS: {args.epochs} SEED: {args.seed}")
print(f" HUMAN-EVAL FULL: base={base_corr}/{base_total} trained={tr_corr}/{tr_total} Δ={tr_corr-base_corr:+d}")
print(f" time: {time.time()-T0:.0f}s")
print("=" * 70)
if __name__ == "__main__":
main()

322
recipe/curriculum_code.py Normal file
View file

@ -0,0 +1,322 @@
"""TinyForge-Zero on CODE with self-difficulty curriculum.
Loop:
1. Generate problem (seeded fresh or amplified/simplified from pool)
2. Greedy solve. Verify against tests.
- If correct → easy → amplify
- If wrong → try 4 sampled attempts
- If at-edge (some pass, some fail) → MINE pair
- If all fail → too hard → simplify
3. Train periodically. Eval on HumanEval.
"""
import os, sys, json, time, re, gc, subprocess, tempfile, argparse, random
os.environ.setdefault("HF_HOME", "/workspace/hf")
os.environ.setdefault("CUDA_VISIBLE_DEVICES", "0")
os.environ["TRANSFORMERS_VERBOSITY"] = "error"
os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
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 run_python(code, timeout=8):
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.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 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
SEED_GEN_PROMPT = """Generate ONE simple Python coding problem with a clear function spec and 3 test assertions.
Output exactly:
```python
def {function_name}({args}):
\"\"\"{description}\"\"\"
{implementation}
# tests
assert {function_name}(...) == ...
assert {function_name}(...) == ...
assert {function_name}(...) == ...
```
Output ONLY the code block."""
AMPLIFY_PROMPT = """Take this Python coding problem and make it HARDER (add an edge case, additional constraint, or trickier logic). Keep the format with function + 3 assert tests.
Original:
```python
{original}
```
Output the harder version (function + tests) in one ```python block."""
SIMPLIFY_PROMPT = """Take this Python coding problem and make it EASIER (remove an edge case, simplify the logic). Keep the format with function + 3 assert tests.
Original:
```python
{original}
```
Output the easier version (function + tests) in one ```python block."""
def parse_problem(text):
code = extract_code(text) if "```" in text else text.strip()
if "def " not in code: return None
lines = code.split("\n")
func_start = next((i for i, l in enumerate(lines) if l.startswith("def ")), None)
if func_start is None: return None
tests = []
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
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
m = re.match(r"def\s+(\w+)\s*\(", lines[func_start])
if not m: return None
fn_name = m.group(1)
sig_lines = []
for i in range(func_start, def_end):
sig_lines.append(lines[i])
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
if i > func_start and '"""' in lines[i] and (i > func_start+1 and '"""' in lines[i-1] or lines[i].count('"""') >= 2):
break
return {"fn_name": fn_name, "signature": "\n".join(sig_lines), "tests": tests,
"canonical": full_solution, "raw": code}
def humaneval_full(model, tok, n=164):
he = list(load_dataset("openai_humaneval", split="test"))[:n]
log(f" HumanEval ({len(he)} problems)")
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
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, len(he)
def make_train_example(r, tok):
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}
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--model", default="Qwen/Qwen2.5-7B")
ap.add_argument("--iterations", type=int, default=16)
ap.add_argument("--problems_per_iter", type=int, default=8)
ap.add_argument("--train_every", type=int, default=4)
ap.add_argument("--seed", type=int, default=42)
ap.add_argument("--tag", required=True)
args = ap.parse_args()
random.seed(args.seed); torch.manual_seed(args.seed)
out_dir = f"/workspace/curriculum_code/{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, dtype=torch.bfloat16, device_map="cuda:0")
log(f" loaded mem={torch.cuda.memory_allocated('cuda:0')/1e9:.1f}GB")
model.eval()
log("INITIAL eval on HumanEval")
base_correct, base_total = humaneval_full(model, tok)
log(f" base: {base_correct}/{base_total}")
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)
accumulated = []
problem_pool = []
for it in range(1, args.iterations + 1):
it_t = time.time()
if not problem_pool:
gen_prompts = [SEED_GEN_PROMPT for _ in range(args.problems_per_iter)]
raw = gen_batch(model, tok, gen_prompts, max_new=400, temperature=0.9)
seeded = []
for r in raw:
parsed = parse_problem(r)
if not parsed: continue
full = parsed["canonical"] + "\n\n" + "\n".join(parsed["tests"])
ok, _ = run_python(full)
if ok: seeded.append(parsed)
problem_pool.extend(seeded)
log(f"iter {it}: seeded {len(seeded)} fresh (pool={len(problem_pool)})")
random.shuffle(problem_pool)
attempt_problems = problem_pool[:args.problems_per_iter]
problem_pool = problem_pool[args.problems_per_iter:]
if not attempt_problems:
log(f"iter {it}: empty pool"); continue
# Greedy solve
greedy_prompts = [f"Implement: {p['signature']}\n\nTests:\n{chr(10).join(p['tests'])}\n\nOutput only the function in one ```python block." for p in attempt_problems]
greedy_outs = gen_batch(model, tok, greedy_prompts, max_new=300, temperature=0.0)
new_pairs = 0
amp_targets = []; sim_targets = []
for p, raw in zip(attempt_problems, greedy_outs):
code = extract_code(raw) if "```" in raw else raw
ok, _ = run_python(code + "\n\n" + "\n".join(p["tests"]))
if ok:
amp_targets.append(p)
else:
# at-edge check via sampling
solve_prompt = f"Implement: {p['signature']}\n\nTests:\n{chr(10).join(p['tests'])}\n\nOutput only the function in one ```python block."
atts = gen_batch(model, tok, [solve_prompt]*4, max_new=300, temperature=0.7)
broken = None; broken_err = None; fixed = None
for ra in atts:
c = extract_code(ra) if "```" in ra else ra
ok2, err = run_python(c + "\n\n" + "\n".join(p["tests"]))
if ok2 and fixed is None: fixed = c
elif not ok2 and broken is None: broken = c; broken_err = err
if broken and fixed: break
if broken and fixed:
accumulated.append({"signature": p["signature"], "tests": p["tests"],
"broken": broken, "error": broken_err, "fixed": fixed})
new_pairs += 1
else:
sim_targets.append(p)
log(f"iter {it}: {len(attempt_problems)} attempted, +{new_pairs} pairs (total: {len(accumulated)}). amp={len(amp_targets)}, sim={len(sim_targets)} [{time.time()-it_t:.0f}s]")
# Generate amplified / simplified for next iter
if amp_targets:
amp_prompts = [AMPLIFY_PROMPT.format(original=p["raw"]) for p in amp_targets[:args.problems_per_iter]]
amp_outs = gen_batch(model, tok, amp_prompts, max_new=400, temperature=0.7)
for r in amp_outs:
parsed = parse_problem(r)
if not parsed: continue
full = parsed["canonical"] + "\n\n" + "\n".join(parsed["tests"])
ok, _ = run_python(full)
if ok: problem_pool.append(parsed)
if sim_targets:
sim_prompts = [SIMPLIFY_PROMPT.format(original=p["raw"]) for p in sim_targets[:args.problems_per_iter//2]]
sim_outs = gen_batch(model, tok, sim_prompts, max_new=400, temperature=0.7)
for r in sim_outs:
parsed = parse_problem(r)
if not parsed: continue
full = parsed["canonical"] + "\n\n" + "\n".join(parsed["tests"])
ok, _ = run_python(full)
if ok: problem_pool.append(parsed)
with open(f"{out_dir}/pairs.jsonl", "w") as fh:
for r in accumulated: fh.write(json.dumps(r) + "\n")
if it % args.train_every == 0 and len(accumulated) >= 10:
log(f" TRAINING on {len(accumulated)} pairs")
tok.padding_side = "right"
ds = HFDataset.from_list([make_train_example(r, tok) for r in accumulated])
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()
tok.padding_side = "left"
model.eval()
corr, tot = humaneval_full(model, tok)
log(f" HumanEval @ iter {it}: {corr}/{tot} Δ={corr-base_correct:+d}")
model.train()
model.eval()
final_correct, final_total = humaneval_full(model, tok)
result = {
"model": args.model, "iterations": args.iterations,
"n_pairs": len(accumulated),
"base": [base_correct, base_total],
"trained": [final_correct, final_total],
"delta": final_correct - base_correct,
"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" CURRICULUM TINYFORGE-ZERO-CODE — {args.model}")
print(f" HumanEval: base={base_correct}/{base_total} trained={final_correct}/{final_total} Δ={final_correct-base_correct:+d}")
print(f" Self-mined pairs: {len(accumulated)}")
print(f" Time: {time.time()-T0:.0f}s")
print("=" * 70)
if __name__ == "__main__":
main()

115
recipe/eval_plus.py
View file

@ -1,115 +0,0 @@
"""Eval our best 14B adapter on HumanEval+ (contamination-resistant hidden tests)."""
import os, json, time, re, subprocess, tempfile, argparse
os.environ.setdefault("HF_HOME", "/workspace/hf")
os.environ["TRANSFORMERS_VERBOSITY"] = "error"
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from datasets import load_dataset
from peft import PeftModel
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=15):
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 gen_batch(model, tok, prompts, max_new=400, 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 a Python coder. Output one ```python block only."},
{"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=False, 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 main():
ap = argparse.ArgumentParser()
ap.add_argument("--model", default="Qwen/Qwen2.5-14B")
ap.add_argument("--adapter", default="/workspace/multi_pair/multi_v1/adapter")
ap.add_argument("--tag", required=True)
args = ap.parse_args()
out_dir = f"/workspace/eval_plus/{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")
if args.adapter and os.path.exists(args.adapter):
log(f" loading adapter from {args.adapter}")
model = PeftModel.from_pretrained(model, args.adapter)
else:
log(" no adapter — base only")
model.eval()
# Load HumanEval+ via evalplus dataset
log("loading HumanEvalPlus dataset")
ds = list(load_dataset("evalplus/humanevalplus", split="test"))
log(f" {len(ds)} problems")
# Eval
log("eval...")
prompts = [p["prompt"] + "\n# Complete the function above." for p in ds]
outs = gen_batch(model, tok, prompts, max_new=400, batch=4)
base_pass, plus_pass = 0, 0
for i, (p, raw) in enumerate(zip(ds, outs)):
code = extract_code(raw) if "```" in raw else raw
full = p["prompt"] + "\n" + code if "def " not in code else code
# Public tests
base_test = full + "\n\n" + p["test"] + f"\n\ncheck({p['entry_point']})"
b = run_python(base_test, timeout=15)
# Plus tests (hidden harder)
plus_check = p.get("plus_input", None)
if plus_check is not None and "plus_test" in p:
plus_test = full + "\n\n" + p["plus_test"] + f"\n\ncheck({p['entry_point']})"
pp = run_python(plus_test, timeout=15)
else:
pp = b # fallback
if b: base_pass += 1
if pp: plus_pass += 1
if (i+1) % 20 == 0:
log(f" {i+1}/{len(ds)}: base={base_pass}, plus={plus_pass}")
result = {"model": args.model, "adapter": args.adapter,
"base_pass": base_pass, "plus_pass": plus_pass, "n": len(ds),
"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" HumanEval+ public: {base_pass}/{len(ds)} plus(hidden): {plus_pass}/{len(ds)}")
print(f" Time: {time.time()-T0:.0f}s")
print("=" * 70)
if __name__ == "__main__":
main()

216
recipe/eval_raw.py
View file

@ -1,216 +0,0 @@
"""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()

283
recipe/math_bootstrap.py Normal file
View file

@ -0,0 +1,283 @@
"""TinyForge-Zero on math word problems.
Same recipe as code bootstrap, different verifier:
- Model generates (word_problem, python_expression_for_answer) pairs
- Python eval gives the canonical numerical answer
- Solver gets word problem only, must produce a number
- Compare solver's number to canonical → broken/fixed pairs
- Train on accumulated pairs
- Eval on GSM8K (held-out)
"""
import os, sys, json, time, re, gc, subprocess, tempfile, argparse, random
os.environ.setdefault("HF_HOME", "/workspace/hf")
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
os.environ["TRANSFORMERS_VERBOSITY"] = "error"
os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
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 safe_eval(expr: str):
"""Eval a numeric Python expression. Returns float or None."""
try:
# Restrict to math operations
allowed = "0123456789+-*/.()% "
if not all(c in allowed or c.isspace() for c in expr): return None
return float(eval(expr, {"__builtins__": {}}, {}))
except Exception:
return None
def extract_answer(text: str):
"""Pull a numeric answer from model output. Looks for last number or boxed."""
# GSM8K style: "#### 42"
m = re.search(r"####\s*(-?\d+(?:\.\d+)?)", text)
if m: return float(m.group(1))
# \boxed{42}
m = re.search(r"\\boxed\{(-?\d+(?:\.\d+)?)\}", text)
if m: return float(m.group(1))
# "answer is 42" or "= 42"
matches = re.findall(r"-?\d+(?:\.\d+)?", text)
if matches:
try: return float(matches[-1])
except: return None
return None
def gen_batch(model, tok, prompts, max_new=400, 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 a careful math tutor."},
{"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 math word problem with a numerical answer. Output exactly this format:
PROBLEM: <a clear word problem with concrete numbers>
EXPRESSION: <a single Python arithmetic expression that evaluates to the answer, e.g. (5*3)+12>
ANSWER: <the numerical answer>
Make the problem grade-school to middle-school level. The expression must evaluate to the answer."""
def parse_generated_problem(text: str):
"""Extract (problem, expression, answer) from model output."""
p_m = re.search(r"PROBLEM:\s*(.+?)(?:\n|EXPRESSION:)", text, re.DOTALL)
e_m = re.search(r"EXPRESSION:\s*(.+?)(?:\n|ANSWER:)", text, re.DOTALL)
a_m = re.search(r"ANSWER:\s*(-?\d+(?:\.\d+)?)", text)
if not (p_m and e_m and a_m): return None
problem = p_m.group(1).strip()
expression = e_m.group(1).strip()
try:
claimed = float(a_m.group(1))
except: return None
if len(problem) < 10 or len(expression) < 1: return None
# Verify: expression evaluates to claimed answer
actual = safe_eval(expression)
if actual is None: return None
if abs(actual - claimed) > 0.01: return None
return {"problem": problem, "expression": expression, "answer": claimed}
SOLVE_PROMPT_TEMPLATE = """Solve this math problem step by step. End with the answer on a new line as: #### <number>
Problem: {problem}"""
def solve_and_check(model, tok, problem_text: str, gold_answer: float, n_attempts: int = 4, temperature: float = 0.7):
"""Sample N attempts, return list of (text, predicted_num, ok)."""
prompt = SOLVE_PROMPT_TEMPLATE.format(problem=problem_text)
outs = gen_batch(model, tok, [prompt] * n_attempts, max_new=400, temperature=temperature)
results = []
for raw in outs:
pred = extract_answer(raw)
ok = pred is not None and abs(pred - gold_answer) < 0.01
results.append({"text": raw, "pred": pred, "ok": ok})
return results
def gsm8k_eval(model, tok, n=200):
ds = list(load_dataset("openai/gsm8k", "main", split="test"))
ds = ds[:n]
log(f" eval on GSM8K ({len(ds)} problems)")
prompts = [SOLVE_PROMPT_TEMPLATE.format(problem=p["question"]) for p in ds]
outs = gen_batch(model, tok, prompts, max_new=400, temperature=0.0, batch=4)
correct = 0
for p, raw in zip(ds, outs):
# GSM8K's answer field has format "step-by-step\n#### 42"
gold_m = re.search(r"####\s*(-?\d+(?:,\d+)*(?:\.\d+)?)", p["answer"])
if not gold_m: continue
gold = float(gold_m.group(1).replace(",", ""))
pred = extract_answer(raw)
if pred is not None and abs(pred - gold) < 0.01: correct += 1
return correct, len(ds)
def make_train_example(r, tok):
user = SOLVE_PROMPT_TEMPLATE.format(problem=r["problem"]) + f"\n\nMy attempt:\n{r['broken']}\n\nThis is wrong. Solve it correctly and end with #### <number>."
assistant = r["fixed"]
msgs_pre = [{"role": "system", "content": "You are a careful math tutor."},
{"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", default="Qwen/Qwen2.5-7B")
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=8)
ap.add_argument("--eval_every", type=int, default=8)
ap.add_argument("--n_eval", type=int, default=200)
ap.add_argument("--seed", type=int, default=42)
ap.add_argument("--tag", required=True)
args = ap.parse_args()
random.seed(args.seed); torch.manual_seed(args.seed)
out_dir = f"/workspace/math/{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"
device = "cuda:0" # CUDA_VISIBLE_DEVICES=1 makes physical GPU 1 appear as cuda:0
model = AutoModelForCausalLM.from_pretrained(args.model, dtype=torch.bfloat16, device_map=device)
log(f" loaded mem={torch.cuda.memory_allocated(device)/1e9:.1f}GB")
# Initial eval
model.eval()
log("INITIAL eval on GSM8K")
init_correct, init_total = gsm8k_eval(model, tok, n=args.n_eval)
log(f" GSM8K base: {init_correct}/{init_total}")
# LoRA
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 problems
gen_prompts = [PROBLEM_GEN_PROMPT for _ in range(args.problems_per_iter)]
raw_problems = gen_batch(model, tok, gen_prompts, max_new=300, temperature=0.9)
# 2. Parse & verify (Python eval of expression)
valid = []
for raw in raw_problems:
parsed = parse_generated_problem(raw)
if parsed: valid.append(parsed)
if not valid:
log(f"iter {it}: 0 valid problems")
iter_stats.append({"iter": it, "valid": 0, "pairs": 0})
continue
# 3. Mine pairs from sampled solver outputs
new_pairs = 0
for p in valid:
attempts = solve_and_check(model, tok, p["problem"], p["answer"], n_attempts=4, temperature=0.7)
ok_atts = [a for a in attempts if a["ok"]]
bad_atts = [a for a in attempts if not a["ok"]]
if ok_atts and bad_atts:
accumulated_pairs.append({
"problem": p["problem"],
"answer": p["answer"],
"broken": bad_atts[0]["text"],
"fixed": ok_atts[0]["text"],
})
new_pairs += 1
log(f"iter {it}: {len(valid)} valid problems, {new_pairs} pairs harvested (total: {len(accumulated_pairs)}) [{time.time()-it_t:.0f}s]")
iter_stats.append({"iter": it, "valid": len(valid), "pairs": new_pairs, "elapsed": time.time()-it_t})
# Save incrementally
with open(f"{out_dir}/pairs.jsonl", "w") as fh:
for r in accumulated_pairs: fh.write(json.dumps(r) + "\n")
# 4. Train every N
if it % args.train_every == 0 and len(accumulated_pairs) >= 10:
log(f" TRAINING on {len(accumulated_pairs)} pairs")
tok.padding_side = "right"
ds = HFDataset.from_list([make_train_example(r, tok) for r in accumulated_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()
tok.padding_side = "left"
# 5. Eval every N
if it % args.eval_every == 0:
model.eval()
corr, tot = gsm8k_eval(model, tok, n=args.n_eval)
log(f" GSM8K @ iter {it}: {corr}/{tot}")
eval_log.append({"iter": it, "correct": corr, "total": tot})
model.train()
# Final eval
model.eval()
final_correct, final_total = gsm8k_eval(model, tok, n=args.n_eval)
eval_log.append({"iter": args.iterations, "correct": final_correct, "total": final_total, "final": True})
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)
print()
print("=" * 70)
print(f" TINYFORGE-ZERO ON MATH ({args.model})")
print(f" GSM8K-mini ({final_total}): base={init_correct} final={final_correct} Δ={final_correct-init_correct:+d}")
print(f" Total pairs mined: {len(accumulated_pairs)}")
print(f" Time: {time.time()-T0:.0f}s")
print("=" * 70)
if __name__ == "__main__":
main()