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:
Rana Usman
parent
c867697f7c
commit
826f934d2e
27 changed files with 4467 additions and 134 deletions
165
evals/confirm.py
Normal file
165
evals/confirm.py
Normal file
|
|
@ -0,0 +1,165 @@
|
|||
"""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()
|
||||
115
evals/eval_plus.py
Normal file
115
evals/eval_plus.py
Normal file
|
|
@ -0,0 +1,115 @@
|
|||
"""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
evals/eval_raw.py
Normal file
216
evals/eval_raw.py
Normal file
|
|
@ -0,0 +1,216 @@
|
|||
"""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()
|
||||
Loading…
Add table
Add a link
Reference in a new issue