mirror of
https://github.com/ranausmanai/tinyforge-zero.git
synced 2026-06-08 20:55:13 +02:00
Ship every paper-referenced experiment script
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:
parent
c867697f7c
commit
826f934d2e
27 changed files with 4467 additions and 134 deletions
103
tts/tts_aime.py
Normal file
103
tts/tts_aime.py
Normal file
|
|
@ -0,0 +1,103 @@
|
|||
"""TTS on AIME (Olympiad math). 90 problems, integer answers 0-999.
|
||||
If 8B+best-of-N hits 30%+, that's matching frontier reasoning models."""
|
||||
import os, json, time, re, argparse
|
||||
os.environ.setdefault("HF_HOME", "/workspace/hf")
|
||||
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_int(text):
|
||||
"""AIME answers are integers 0-999. Try \boxed first, fall back to last integer."""
|
||||
m = re.search(r"\\boxed\{(\d+)\}", text)
|
||||
if m:
|
||||
try: return int(m.group(1))
|
||||
except: return None
|
||||
# Last integer in last few lines
|
||||
lines = text.strip().split("\n")
|
||||
for line in reversed(lines[-5:]):
|
||||
nums = re.findall(r"\b(\d+)\b", line)
|
||||
if nums:
|
||||
try: return int(nums[-1])
|
||||
except: pass
|
||||
return None
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--model", required=True)
|
||||
ap.add_argument("--n_samples", type=int, default=8)
|
||||
ap.add_argument("--temperature", type=float, default=0.7)
|
||||
ap.add_argument("--tag", required=True)
|
||||
args = ap.parse_args()
|
||||
|
||||
out_dir = f"/workspace/tts_aime/{args.tag}"
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
|
||||
from vllm import LLM, SamplingParams
|
||||
from transformers import AutoTokenizer
|
||||
log(f"loading {args.model}")
|
||||
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.90, max_model_len=3072)
|
||||
log(f" loaded")
|
||||
|
||||
ds = list(load_dataset("AI-MO/aimo-validation-aime", split="train"))
|
||||
log(f" AIME: {len(ds)} problems")
|
||||
|
||||
SYS = "You are a careful math problem solver. AIME answers are integers between 0 and 999. End with \\boxed{integer}."
|
||||
UTMPL = "Solve this AIME problem. Show your reasoning, then put the final integer answer in \\boxed{{...}}.\n\nProblem: {problem}\n\nSolution:"
|
||||
prompts = []
|
||||
for p in ds:
|
||||
msgs = [{"role": "system", "content": SYS},
|
||||
{"role": "user", "content": UTMPL.format(problem=p["problem"])}]
|
||||
try:
|
||||
prompts.append(tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True))
|
||||
except Exception:
|
||||
prompts.append(UTMPL.format(problem=p["problem"]))
|
||||
|
||||
log("=== GREEDY ===")
|
||||
sp_g = SamplingParams(temperature=0, max_tokens=2000)
|
||||
t0 = time.time()
|
||||
g_outs = [o.outputs[0].text for o in llm.generate(prompts, sp_g, use_tqdm=False)]
|
||||
log(f" gen in {time.time()-t0:.1f}s")
|
||||
g_correct = 0
|
||||
for p, raw in zip(ds, g_outs):
|
||||
pred = extract_int(raw)
|
||||
gold = int(p["answer"])
|
||||
if pred == gold: g_correct += 1
|
||||
log(f" GREEDY: {g_correct}/{len(ds)} ({100*g_correct/len(ds):.1f}%)")
|
||||
|
||||
log(f"=== BEST-OF-{args.n_samples} (temp={args.temperature}) ===")
|
||||
sp_s = SamplingParams(temperature=args.temperature, top_p=0.95, max_tokens=2000, n=args.n_samples)
|
||||
t0 = time.time()
|
||||
s_outs = llm.generate(prompts, sp_s, use_tqdm=False)
|
||||
log(f" gen in {time.time()-t0:.1f}s")
|
||||
bN_correct = 0
|
||||
for p, outset in zip(ds, s_outs):
|
||||
gold = int(p["answer"])
|
||||
for o in outset.outputs:
|
||||
pred = extract_int(o.text)
|
||||
if pred == gold:
|
||||
bN_correct += 1; break
|
||||
|
||||
result = {"model": args.model, "n_samples": args.n_samples, "temperature": args.temperature,
|
||||
"greedy": g_correct, "best_of_N": bN_correct, "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" {args.model} — AIME ({len(ds)} problems)")
|
||||
print(f" Greedy: {g_correct}/{len(ds)} ({100*g_correct/len(ds):.1f}%)")
|
||||
print(f" Best-of-{args.n_samples}: {bN_correct}/{len(ds)} ({100*bN_correct/len(ds):.1f}%)")
|
||||
print(f" TTS Lift: +{bN_correct - g_correct} ({100*(bN_correct-g_correct)/len(ds):.1f}pp)")
|
||||
print(f" Time: {time.time()-T0:.0f}s")
|
||||
print("=" * 70)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
126
tts/tts_humaneval.py
Normal file
126
tts/tts_humaneval.py
Normal file
|
|
@ -0,0 +1,126 @@
|
|||
"""TTS on HumanEval+ (contamination-resistant) to verify the 92% isn't memorization."""
|
||||
import os, json, time, subprocess, tempfile, argparse
|
||||
os.environ.setdefault("HF_HOME", "/workspace/hf")
|
||||
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=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 main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--model", required=True)
|
||||
ap.add_argument("--n_samples", type=int, default=8)
|
||||
ap.add_argument("--temperature", type=float, default=0.6)
|
||||
ap.add_argument("--tag", required=True)
|
||||
args = ap.parse_args()
|
||||
|
||||
out_dir = f"/workspace/tts_hep/{args.tag}"
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
|
||||
from vllm import LLM, SamplingParams
|
||||
from transformers import AutoTokenizer
|
||||
log(f"loading {args.model}")
|
||||
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.90, max_model_len=2048)
|
||||
log(f" loaded")
|
||||
|
||||
hep = list(load_dataset("evalplus/humanevalplus", split="test"))
|
||||
log(f" HE+: {len(hep)} problems")
|
||||
|
||||
prompts = []
|
||||
for p in hep:
|
||||
try:
|
||||
msgs = [{"role": "system", "content": "You are a Python coder. Output one ```python block only."},
|
||||
{"role": "user", "content": p["prompt"] + "\n# Complete the function above."}]
|
||||
prompts.append(tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True))
|
||||
except Exception:
|
||||
prompts.append(p["prompt"])
|
||||
|
||||
log("=== GREEDY ===")
|
||||
sp_g = SamplingParams(temperature=0, max_tokens=400)
|
||||
g_outs = [o.outputs[0].text for o in llm.generate(prompts, sp_g, use_tqdm=False)]
|
||||
base_pass, plus_pass = 0, 0
|
||||
for p, raw in zip(hep, g_outs):
|
||||
code = extract_code(raw) if "```" in raw else raw
|
||||
full = p["prompt"] + "\n" + code if "def " not in code else code
|
||||
# base test
|
||||
b_test = full + "\n\n" + p["test"] + f"\n\ncheck({p['entry_point']})"
|
||||
b_ok = run_python(b_test, 15)
|
||||
if b_ok: base_pass += 1
|
||||
# plus test (harder, hidden cases)
|
||||
if "plus_test" in p:
|
||||
p_test = full + "\n\n" + p["plus_test"] + f"\n\ncheck({p['entry_point']})"
|
||||
if run_python(p_test, 15): plus_pass += 1
|
||||
else:
|
||||
if b_ok: plus_pass += 1
|
||||
log(f" GREEDY base: {base_pass}/{len(hep)} plus(hidden): {plus_pass}/{len(hep)}")
|
||||
|
||||
log(f"=== BEST-OF-{args.n_samples} (temp={args.temperature}) ===")
|
||||
sp_s = SamplingParams(temperature=args.temperature, top_p=0.95, max_tokens=400, n=args.n_samples)
|
||||
s_outs = llm.generate(prompts, sp_s, use_tqdm=False)
|
||||
bN_base, bN_plus = 0, 0
|
||||
for p, outset in zip(hep, s_outs):
|
||||
attempts = [o.text for o in outset.outputs]
|
||||
base_ok_any = False
|
||||
plus_ok_any = False
|
||||
for a in attempts:
|
||||
code = extract_code(a) if "```" in a else a
|
||||
full = p["prompt"] + "\n" + code if "def " not in code else code
|
||||
b_test = full + "\n\n" + p["test"] + f"\n\ncheck({p['entry_point']})"
|
||||
b_ok = run_python(b_test, 15)
|
||||
if b_ok and not base_ok_any:
|
||||
base_ok_any = True
|
||||
if "plus_test" in p:
|
||||
p_test = full + "\n\n" + p["plus_test"] + f"\n\ncheck({p['entry_point']})"
|
||||
p_ok = run_python(p_test, 15)
|
||||
if p_ok and not plus_ok_any:
|
||||
plus_ok_any = True
|
||||
elif b_ok and not plus_ok_any:
|
||||
plus_ok_any = True
|
||||
if base_ok_any and plus_ok_any: break
|
||||
if base_ok_any: bN_base += 1
|
||||
if plus_ok_any: bN_plus += 1
|
||||
|
||||
result = {"model": args.model, "n_samples": args.n_samples, "temperature": args.temperature,
|
||||
"greedy_base": base_pass, "greedy_plus": plus_pass,
|
||||
"best_of_N_base": bN_base, "best_of_N_plus": bN_plus,
|
||||
"n": len(hep), "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} — HumanEval+ ({len(hep)} problems)")
|
||||
print(f" Greedy base: {base_pass}/{len(hep)} ({100*base_pass/len(hep):.1f}%)")
|
||||
print(f" Greedy plus (hard): {plus_pass}/{len(hep)} ({100*plus_pass/len(hep):.1f}%)")
|
||||
print(f" Best-of-{args.n_samples} base: {bN_base}/{len(hep)} ({100*bN_base/len(hep):.1f}%)")
|
||||
print(f" Best-of-{args.n_samples} plus: {bN_plus}/{len(hep)} ({100*bN_plus/len(hep):.1f}%)")
|
||||
print(f" Time: {time.time()-T0:.0f}s")
|
||||
print("=" * 70)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
125
tts/tts_math500.py
Normal file
125
tts/tts_math500.py
Normal file
|
|
@ -0,0 +1,125 @@
|
|||
"""TTS on MATH-500: greedy + best-of-N pass@1.
|
||||
|
||||
If TTS works on math like it does on code, we should see major lift.
|
||||
"""
|
||||
import os, json, time, re, argparse
|
||||
os.environ.setdefault("HF_HOME", "/workspace/hf")
|
||||
os.environ["TRANSFORMERS_VERBOSITY"] = "error"
|
||||
|
||||
import torch
|
||||
from datasets import load_dataset
|
||||
import sympy
|
||||
from sympy.parsing.latex import parse_latex
|
||||
|
||||
T0 = time.time()
|
||||
def log(m): print(f"[{time.time()-T0:7.1f}s] {m}", flush=True)
|
||||
|
||||
|
||||
def extract_boxed(text):
|
||||
idx = text.rfind("\\boxed{")
|
||||
if idx < 0: return None
|
||||
start = idx + len("\\boxed{")
|
||||
depth = 1; i = start
|
||||
while i < len(text) and depth > 0:
|
||||
if text[i] == "{": depth += 1
|
||||
elif text[i] == "}": depth -= 1
|
||||
i += 1
|
||||
if depth != 0: return None
|
||||
return text[start:i-1].strip()
|
||||
|
||||
|
||||
def normalize(s):
|
||||
if s is None: return None
|
||||
s = s.strip()
|
||||
s = re.sub(r"^\$|\$$", "", s).strip()
|
||||
s = re.sub(r"\\text\{([^}]*)\}", r"\1", s)
|
||||
s = re.sub(r"\\mbox\{([^}]*)\}", r"\1", s)
|
||||
s = re.sub(r"(?<=\d),(?=\d)", "", s)
|
||||
s = s.replace("\\left", "").replace("\\right", "").replace("^\\circ", "").replace("^{\\circ}", "")
|
||||
return s.strip()
|
||||
|
||||
|
||||
def sympy_equal(a, b):
|
||||
if a is None or b is None: return False
|
||||
a, b = normalize(a), normalize(b)
|
||||
if a == b: return True
|
||||
try:
|
||||
ea = parse_latex(a); eb = parse_latex(b)
|
||||
if sympy.simplify(ea - eb) == 0: return True
|
||||
except Exception: pass
|
||||
try:
|
||||
if abs(float(a) - float(b)) < 1e-6: return True
|
||||
except Exception: pass
|
||||
return False
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--model", required=True)
|
||||
ap.add_argument("--n_samples", type=int, default=8)
|
||||
ap.add_argument("--temperature", type=float, default=0.7)
|
||||
ap.add_argument("--tag", required=True)
|
||||
args = ap.parse_args()
|
||||
|
||||
out_dir = f"/workspace/tts_math/{args.tag}"
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
|
||||
from vllm import LLM, SamplingParams
|
||||
from transformers import AutoTokenizer
|
||||
log(f"loading {args.model}")
|
||||
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.90, max_model_len=2048)
|
||||
log(f" loaded")
|
||||
|
||||
ds = list(load_dataset("HuggingFaceH4/MATH-500", split="test"))
|
||||
log(f" MATH-500: {len(ds)} problems")
|
||||
|
||||
SYS = "You are a careful math problem solver. End with \\boxed{answer}."
|
||||
USER_TEMPLATE = "Solve this competition math problem. Show your reasoning, then put the final answer in \\boxed{{...}}.\n\nProblem: {problem}\n\nSolution:"
|
||||
prompts = []
|
||||
for p in ds:
|
||||
msgs = [{"role": "system", "content": SYS},
|
||||
{"role": "user", "content": USER_TEMPLATE.format(problem=p["problem"])}]
|
||||
try:
|
||||
prompts.append(tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True))
|
||||
except Exception:
|
||||
prompts.append(USER_TEMPLATE.format(problem=p["problem"]))
|
||||
|
||||
# Greedy
|
||||
log("=== GREEDY ===")
|
||||
sp_g = SamplingParams(temperature=0, max_tokens=800)
|
||||
t0 = time.time()
|
||||
g_outs = [o.outputs[0].text for o in llm.generate(prompts, sp_g, use_tqdm=False)]
|
||||
log(f" gen in {time.time()-t0:.1f}s")
|
||||
g_correct = sum(1 for p, raw in zip(ds, g_outs) if sympy_equal(extract_boxed(raw), p["answer"]))
|
||||
log(f" GREEDY: {g_correct}/{len(ds)} ({100*g_correct/len(ds):.1f}%)")
|
||||
|
||||
# Best-of-N (any correct)
|
||||
log(f"=== BEST-OF-{args.n_samples} (temp={args.temperature}) ===")
|
||||
sp_s = SamplingParams(temperature=args.temperature, top_p=0.95, max_tokens=800, n=args.n_samples)
|
||||
t0 = time.time()
|
||||
s_outs = llm.generate(prompts, sp_s, use_tqdm=False)
|
||||
log(f" gen in {time.time()-t0:.1f}s")
|
||||
bN_correct = 0
|
||||
for p, outset in zip(ds, s_outs):
|
||||
for o in outset.outputs:
|
||||
if sympy_equal(extract_boxed(o.text), p["answer"]):
|
||||
bN_correct += 1; break
|
||||
|
||||
result = {"model": args.model, "n_samples": args.n_samples, "temperature": args.temperature,
|
||||
"greedy": g_correct, "best_of_N": bN_correct, "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" {args.model} — MATH-500 ({len(ds)} problems)")
|
||||
print(f" Greedy: {g_correct}/{len(ds)} ({100*g_correct/len(ds):.1f}%)")
|
||||
print(f" Best-of-{args.n_samples}: {bN_correct}/{len(ds)} ({100*bN_correct/len(ds):.1f}%)")
|
||||
print(f" TTS Lift: +{bN_correct - g_correct} ({100*(bN_correct-g_correct)/len(ds):.1f}pp)")
|
||||
print(f" Time: {time.time()-T0:.0f}s")
|
||||
print("=" * 70)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
135
tts/tts_qwen14b_recipe.py
Normal file
135
tts/tts_qwen14b_recipe.py
Normal file
|
|
@ -0,0 +1,135 @@
|
|||
"""Test-time scaling on Qwen2.5-14B-Base + multi_v1 adapter.
|
||||
|
||||
For each HumanEval problem:
|
||||
1. Sample 8 attempts at temp=0.6 from the trained model.
|
||||
2. Run each attempt against the tests.
|
||||
3. Accept the first that passes → pass@1 with best-of-N selection.
|
||||
|
||||
Compared to greedy pass@1 (which gave 80.5%), this should push higher.
|
||||
"""
|
||||
import os, json, time, re, subprocess, tempfile, argparse
|
||||
os.environ.setdefault("HF_HOME", "/workspace/hf")
|
||||
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=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 main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--model", default="Qwen/Qwen2.5-14B")
|
||||
ap.add_argument("--adapter", default="/workspace/multi_v1_adapter")
|
||||
ap.add_argument("--n_samples", type=int, default=8)
|
||||
ap.add_argument("--temperature", type=float, default=0.6)
|
||||
ap.add_argument("--tag", required=True)
|
||||
args = ap.parse_args()
|
||||
|
||||
out_dir = f"/workspace/tts/{args.tag}"
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
|
||||
from vllm import LLM
|
||||
from vllm.lora.request import LoRARequest
|
||||
from transformers import AutoTokenizer
|
||||
log(f"loading {args.model} with adapter {args.adapter}")
|
||||
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.90, max_model_len=2048,
|
||||
enable_lora=True, max_lora_rank=32)
|
||||
lora_req = LoRARequest("multi_v1", 1, args.adapter)
|
||||
log(f" loaded")
|
||||
|
||||
he = list(load_dataset("openai_humaneval", split="test"))
|
||||
log(f" HE: {len(he)} problems")
|
||||
|
||||
# --- Greedy baseline (with adapter)
|
||||
log("=== GREEDY pass@1 (with adapter) ===")
|
||||
from vllm import SamplingParams
|
||||
sp_greedy = SamplingParams(temperature=0, max_tokens=400)
|
||||
# Use chat template for Qwen2.5 (it has one)
|
||||
prompts = []
|
||||
for p in he:
|
||||
msgs = [{"role": "system", "content": "You are a Python coder. Output one ```python block only."},
|
||||
{"role": "user", "content": p["prompt"] + "\n# Complete the function above."}]
|
||||
prompts.append(tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True))
|
||||
t0 = time.time()
|
||||
greedy_outs = [o.outputs[0].text for o in llm.generate(prompts, sp_greedy, lora_request=lora_req, use_tqdm=False)]
|
||||
log(f" greedy gen in {time.time()-t0:.1f}s")
|
||||
greedy_correct = 0
|
||||
for p, raw in zip(he, greedy_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, 15): greedy_correct += 1
|
||||
log(f" GREEDY pass@1: {greedy_correct}/{len(he)} ({100*greedy_correct/len(he):.1f}%)")
|
||||
|
||||
# --- Test-time scaling: sample N, take first that passes (best-of-N pass@1)
|
||||
log(f"=== TEST-TIME SCALING: N={args.n_samples}, temp={args.temperature} ===")
|
||||
sp_sample = SamplingParams(temperature=args.temperature, top_p=0.95,
|
||||
max_tokens=400, n=args.n_samples)
|
||||
t0 = time.time()
|
||||
sample_outs = llm.generate(prompts, sp_sample, lora_request=lora_req, use_tqdm=False)
|
||||
log(f" sampling gen in {time.time()-t0:.1f}s")
|
||||
|
||||
t1 = time.time()
|
||||
bestN_correct = 0
|
||||
per_problem = []
|
||||
for p, outset in zip(he, sample_outs):
|
||||
attempts = [o.text for o in outset.outputs]
|
||||
any_pass = False
|
||||
for a in attempts:
|
||||
code = extract_code(a) if "```" in a else a
|
||||
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, 15):
|
||||
any_pass = True
|
||||
break
|
||||
if any_pass: bestN_correct += 1
|
||||
per_problem.append({"task_id": p["task_id"], "best_of_N_pass": any_pass})
|
||||
log(f" verify done in {time.time()-t1:.1f}s")
|
||||
|
||||
result = {
|
||||
"model": args.model, "adapter": args.adapter,
|
||||
"n_samples": args.n_samples, "temperature": args.temperature,
|
||||
"greedy_passN": greedy_correct,
|
||||
"best_of_N_passN": bestN_correct,
|
||||
"n_total": len(he),
|
||||
"elapsed_s": time.time()-T0,
|
||||
}
|
||||
with open(f"{out_dir}/result.json", "w") as fh: json.dump(result, fh, indent=2)
|
||||
with open(f"{out_dir}/per_problem.json", "w") as fh: json.dump(per_problem, fh, indent=2)
|
||||
|
||||
print()
|
||||
print("=" * 70)
|
||||
print(f" {args.model} + adapter {args.adapter}")
|
||||
print(f" HumanEval:")
|
||||
print(f" Greedy pass@1: {greedy_correct}/{len(he)} ({100*greedy_correct/len(he):.1f}%)")
|
||||
print(f" Best-of-{args.n_samples} pass@1: {bestN_correct}/{len(he)} ({100*bestN_correct/len(he):.1f}%)")
|
||||
print(f" Lift: +{bestN_correct - greedy_correct} ({100*(bestN_correct-greedy_correct)/len(he):.1f}pp)")
|
||||
print(f" Time: {time.time()-T0:.0f}s")
|
||||
print("=" * 70)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
118
tts/tts_qwen3_8b_raw_control.py
Normal file
118
tts/tts_qwen3_8b_raw_control.py
Normal file
|
|
@ -0,0 +1,118 @@
|
|||
"""Control: Qwen3-8B-Base RAW (no recipe) + best-of-8 on HumanEval.
|
||||
|
||||
Tells us if the 89.6% headline on 14B+recipe is driven by recipe or by test-time scaling.
|
||||
"""
|
||||
import os, json, time, re, subprocess, tempfile, argparse
|
||||
os.environ.setdefault("HF_HOME", "/workspace/hf")
|
||||
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=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 main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--model", required=True)
|
||||
ap.add_argument("--n_samples", type=int, default=8)
|
||||
ap.add_argument("--temperature", type=float, default=0.6)
|
||||
ap.add_argument("--tag", required=True)
|
||||
args = ap.parse_args()
|
||||
|
||||
out_dir = f"/workspace/tts_raw/{args.tag}"
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
|
||||
from vllm import LLM, SamplingParams
|
||||
from transformers import AutoTokenizer
|
||||
log(f"loading {args.model} (no adapter)")
|
||||
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.90, max_model_len=2048)
|
||||
log(f" loaded")
|
||||
|
||||
he = list(load_dataset("openai_humaneval", split="test"))
|
||||
log(f" HE: {len(he)} problems")
|
||||
|
||||
# Try chat-template style if available, else raw
|
||||
prompts = []
|
||||
for p in he:
|
||||
try:
|
||||
msgs = [{"role": "system", "content": "You are a Python coder. Output one ```python block only."},
|
||||
{"role": "user", "content": p["prompt"] + "\n# Complete the function above."}]
|
||||
prompts.append(tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True))
|
||||
except Exception:
|
||||
prompts.append(p["prompt"])
|
||||
|
||||
# --- Greedy
|
||||
log("=== GREEDY pass@1 ===")
|
||||
sp_g = SamplingParams(temperature=0, max_tokens=400)
|
||||
t0 = time.time()
|
||||
g_outs = [o.outputs[0].text for o in llm.generate(prompts, sp_g, use_tqdm=False)]
|
||||
log(f" greedy gen in {time.time()-t0:.1f}s")
|
||||
g_correct = 0
|
||||
for p, raw in zip(he, g_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, 15): g_correct += 1
|
||||
log(f" GREEDY pass@1: {g_correct}/{len(he)} ({100*g_correct/len(he):.1f}%)")
|
||||
|
||||
# --- Best-of-N
|
||||
log(f"=== BEST-OF-{args.n_samples} (temp={args.temperature}) ===")
|
||||
sp_s = SamplingParams(temperature=args.temperature, top_p=0.95, max_tokens=400, n=args.n_samples)
|
||||
t0 = time.time()
|
||||
s_outs = llm.generate(prompts, sp_s, use_tqdm=False)
|
||||
log(f" sampling gen in {time.time()-t0:.1f}s")
|
||||
t1 = time.time()
|
||||
bN_correct = 0
|
||||
for p, outset in zip(he, s_outs):
|
||||
attempts = [o.text for o in outset.outputs]
|
||||
for a in attempts:
|
||||
code = extract_code(a) if "```" in a else a
|
||||
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, 15):
|
||||
bN_correct += 1
|
||||
break
|
||||
log(f" verify in {time.time()-t1:.1f}s")
|
||||
|
||||
result = {
|
||||
"model": args.model, "n_samples": args.n_samples, "temperature": args.temperature,
|
||||
"greedy_passN": g_correct, "best_of_N_passN": bN_correct, "n_total": len(he),
|
||||
"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} (NO ADAPTER) — HumanEval")
|
||||
print(f" Greedy pass@1: {g_correct}/{len(he)} ({100*g_correct/len(he):.1f}%)")
|
||||
print(f" Best-of-{args.n_samples} pass@1: {bN_correct}/{len(he)} ({100*bN_correct/len(he):.1f}%)")
|
||||
print(f" Lift from TTS: +{bN_correct - g_correct} ({100*(bN_correct-g_correct)/len(he):.1f}pp)")
|
||||
print(f" Time: {time.time()-T0:.0f}s")
|
||||
print("=" * 70)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
165
tts/tts_scaling.py
Normal file
165
tts/tts_scaling.py
Normal file
|
|
@ -0,0 +1,165 @@
|
|||
"""TTS scaling sweep: pass@1 across N samples for HE + HE+ + MATH-500."""
|
||||
import os, json, time, re, subprocess, tempfile, argparse
|
||||
os.environ.setdefault("HF_HOME", "/workspace/hf")
|
||||
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(t):
|
||||
if "```python" in t: t = t.split("```python", 1)[1]
|
||||
elif "```" in t: t = t.split("```", 1)[1]
|
||||
if "```" in t: t = t.split("```", 1)[0]
|
||||
return t.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 extract_boxed(text):
|
||||
idx = text.rfind("\\boxed{")
|
||||
if idx < 0: return None
|
||||
start = idx + len("\\boxed{"); depth = 1; i = start
|
||||
while i < len(text) and depth > 0:
|
||||
if text[i] == "{": depth += 1
|
||||
elif text[i] == "}": depth -= 1
|
||||
i += 1
|
||||
if depth != 0: return None
|
||||
return text[start:i-1].strip()
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--model", required=True)
|
||||
ap.add_argument("--tag", required=True)
|
||||
ap.add_argument("--out_dir", required=True)
|
||||
args = ap.parse_args()
|
||||
|
||||
os.makedirs(args.out_dir, exist_ok=True)
|
||||
|
||||
from vllm import LLM, SamplingParams
|
||||
from transformers import AutoTokenizer
|
||||
log(f"loading {args.model}")
|
||||
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("loaded")
|
||||
|
||||
he = list(load_dataset("openai_humaneval", split="test"))
|
||||
math500 = list(load_dataset("HuggingFaceH4/MATH-500", split="test"))[:200]
|
||||
|
||||
# Build prompts
|
||||
he_prompts = []
|
||||
for p in he:
|
||||
try:
|
||||
msgs = [{"role": "system", "content": "You are a Python coder. Output one ```python block only."},
|
||||
{"role": "user", "content": p["prompt"] + "\n# Complete the function above."}]
|
||||
he_prompts.append(tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True))
|
||||
except Exception:
|
||||
he_prompts.append(p["prompt"])
|
||||
|
||||
math_prompts = []
|
||||
UTMPL = "Solve this competition math problem. End with \\boxed{{...}}.\n\nProblem: {p}\n\nSolution:"
|
||||
for p in math500:
|
||||
try:
|
||||
msgs = [{"role": "system", "content": "Math solver. End with \\boxed{answer}."},
|
||||
{"role": "user", "content": UTMPL.format(p=p["problem"])}]
|
||||
math_prompts.append(tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True))
|
||||
except Exception:
|
||||
math_prompts.append(UTMPL.format(p=p["problem"]))
|
||||
|
||||
# Generate max-N samples ONCE per task (N=32), then compute pass@k for k ∈ {1, 2, 4, 8, 16, 32}
|
||||
MAX_N = 32
|
||||
sp = SamplingParams(temperature=0.6, top_p=0.95, max_tokens=600, n=MAX_N)
|
||||
log(f"generating MAX_N={MAX_N} samples per task")
|
||||
t0 = time.time()
|
||||
he_outs = llm.generate(he_prompts, sp, use_tqdm=False)
|
||||
log(f" HE gen in {time.time()-t0:.1f}s")
|
||||
t0 = time.time()
|
||||
math_outs = llm.generate(math_prompts, sp, use_tqdm=False)
|
||||
log(f" MATH gen in {time.time()-t0:.1f}s")
|
||||
|
||||
# Compute correctness for each sample
|
||||
def he_correct(p, raw):
|
||||
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']})"
|
||||
return run_python(test_code, 10)
|
||||
|
||||
log("verifying HE samples...")
|
||||
he_results = [] # per task: list of bool
|
||||
for p, outset in zip(he, he_outs):
|
||||
per_task = []
|
||||
for o in outset.outputs:
|
||||
per_task.append(he_correct(p, o.text))
|
||||
he_results.append(per_task)
|
||||
log(f" HE verify done")
|
||||
|
||||
import sympy
|
||||
from sympy.parsing.latex import parse_latex
|
||||
def sympy_eq(a, b):
|
||||
if a is None or b is None: return False
|
||||
a, b = a.strip(), b.strip()
|
||||
if a == b: return True
|
||||
try:
|
||||
if sympy.simplify(parse_latex(a) - parse_latex(b)) == 0: return True
|
||||
except Exception: pass
|
||||
try:
|
||||
if abs(float(a) - float(b)) < 1e-6: return True
|
||||
except Exception: pass
|
||||
return False
|
||||
|
||||
log("verifying MATH samples...")
|
||||
math_results = []
|
||||
for p, outset in zip(math500, math_outs):
|
||||
per_task = []
|
||||
for o in outset.outputs:
|
||||
pred = extract_boxed(o.text)
|
||||
per_task.append(sympy_eq(pred, p["answer"]))
|
||||
math_results.append(per_task)
|
||||
log(f" MATH verify done")
|
||||
|
||||
# Compute pass@k for each k
|
||||
NS = [1, 2, 4, 8, 16, 32]
|
||||
def best_of_k(results, k):
|
||||
return sum(1 for r in results if any(r[:k]))
|
||||
|
||||
he_scaling = {k: best_of_k(he_results, k) for k in NS}
|
||||
math_scaling = {k: best_of_k(math_results, k) for k in NS}
|
||||
|
||||
result = {
|
||||
"model": args.model, "tag": args.tag, "MAX_N": MAX_N,
|
||||
"humaneval_total": len(he),
|
||||
"math500_total": len(math500),
|
||||
"he_pass_at_k": he_scaling,
|
||||
"math500_pass_at_k": math_scaling,
|
||||
"elapsed_s": time.time() - T0,
|
||||
}
|
||||
with open(f"{args.out_dir}/result.json", "w") as fh: json.dump(result, fh, indent=2)
|
||||
|
||||
print()
|
||||
print("=" * 70)
|
||||
print(f" {args.model} — TTS SCALING SWEEP")
|
||||
print(f" N HE MATH-500")
|
||||
for k in NS:
|
||||
print(f" {k:>3} {he_scaling[k]:>3}/{len(he)} ({100*he_scaling[k]/len(he):.1f}%) "
|
||||
f"{math_scaling[k]:>3}/{len(math500)} ({100*math_scaling[k]/len(math500):.1f}%)")
|
||||
print(f" Time: {time.time()-T0:.0f}s")
|
||||
print("=" * 70)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Loading…
Add table
Add a link
Reference in a new issue