Spaces:
Running
Running
| """DPO retrain via AgentTune's closed-loop retrain path (run inside an HF GPU Job). | |
| Consumes the healer's verified pairs, builds agenttune TrainingExample objects, and | |
| trains an adapter-only LoRA through agenttune's own | |
| examples_to_dpo_dataset -> build_retrainer (create_agentic_trainer) -> train() | |
| exactly as the closed loop's run_retrain does (we call build_retrainer so we can save | |
| the adapter explicitly afterwards). | |
| Train/eval consistency: the student evaluates with the Cohere chat template | |
| (student.build_prompt). examples_to_dpo_dataset flattens prompt messages as | |
| "role: content", BUT with role="" it emits content verbatim - so we pre-apply the | |
| chat template here and hand it the final string. Training sees exactly the text the | |
| model sees at inference. | |
| Launched by launch_job.py, which injects --data/--out/--push-repo/--push-path. | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import json | |
| import os | |
| import sys | |
| def main() -> None: | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--data", default=None) # injected by launcher; unused | |
| ap.add_argument("--out", required=True) | |
| ap.add_argument("--push-repo", default=None) | |
| ap.add_argument("--push-path", default=None) | |
| ap.add_argument("--pairs-repo", default="HideIron/selfheal-sql-runs") | |
| ap.add_argument("--pairs-path", default="heal/baseline-pool600/pairs.jsonl") | |
| ap.add_argument("--model", default="CohereLabs/tiny-aya-global") | |
| ap.add_argument("--adapter-repo", default="HideIron/tiny-aya-global-sql-selfheal") | |
| ap.add_argument("--epochs", type=float, default=3.0) | |
| ap.add_argument("--lr", type=float, default=1e-5) | |
| ap.add_argument("--beta", type=float, default=0.1) | |
| ap.add_argument("--lora-r", type=int, default=16) | |
| ap.add_argument("--batch-size", type=int, default=2) | |
| ap.add_argument("--grad-accum", type=int, default=8) | |
| ap.add_argument("--max-steps", type=int, default=None, help="smoke-test cap") | |
| ap.add_argument("--max-prompt-len", type=int, default=1408) | |
| ap.add_argument("--rpo-alpha", type=float, default=0.0) | |
| ap.add_argument("--algorithm", default="dpo", choices=["dpo", "bco"]) | |
| ap.add_argument("--replay-run", default=None, | |
| help="BCO only: student run whose CORRECT answers are added as " | |
| "positive (label=True) replay rows against forgetting") | |
| ap.add_argument("--no-push-adapter", action="store_true") | |
| args = ap.parse_args() | |
| from huggingface_hub import HfApi, hf_hub_download | |
| from transformers import AutoTokenizer | |
| from agenttune.decide.closed_loop.contracts import TrainingExample | |
| from agenttune.decide.closed_loop.retrain_config import RetrainConfig, build_retrainer | |
| pp = hf_hub_download(args.pairs_repo, args.pairs_path, repo_type="dataset") | |
| pairs = [json.loads(l) for l in open(pp, encoding="utf-8")] | |
| healed = [p for p in pairs if p.get("healed") and p.get("chosen_sql")] | |
| print(f"[dpo] {len(healed)} verified pairs of {len(pairs)} failures", flush=True) | |
| tok = AutoTokenizer.from_pretrained(args.model) | |
| def templ(user_content: str) -> str: | |
| return tok.apply_chat_template( | |
| [{"role": "user", "content": user_content}], | |
| add_generation_prompt=True, tokenize=False, | |
| ) | |
| examples = [ | |
| TrainingExample( | |
| trajectory_id=p["qid"], | |
| original_failure_type=p["status_before"], | |
| root_cause=p["root_cause"], | |
| # role "" -> _messages_to_text emits content verbatim (no "user: " prefix) | |
| prompt=[{"role": "", "content": templ(p["prompt_user_content"])}], | |
| chosen=[{"role": "", "content": p["chosen_sql"]}], | |
| rejected=[{"role": "", "content": p["rejected_sql"]}], | |
| ) | |
| for p in healed | |
| ] | |
| if args.replay_run: | |
| assert args.algorithm == "bco", "success replay needs BCO's independent rows" | |
| import re as _re | |
| def _extract(t): # minimal extract_sql (fences + first statement) | |
| t = _re.sub(r"<think>.*?</think>", "", t, flags=_re.S) | |
| m = _re.search(r"```(?:sql)?\s*(.*?)```", t, _re.S | _re.I) | |
| if m: | |
| t = m.group(1) | |
| m = _re.search(r"\b(select|with)\b", t, _re.I) | |
| return t[m.start():].split(";")[0].strip() if m else t.strip() | |
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) | |
| from student import PROMPT # the exact student prompt builder | |
| split_rows = {} | |
| sp = os.path.join(args.data or "/data", "splits", "pool.jsonl") | |
| for s in map(json.loads, open(sp, encoding="utf-8")): | |
| split_rows[s["qid"]] = s | |
| rp2 = hf_hub_download(args.pairs_repo, f"runs/{args.replay_run}/results.jsonl", | |
| repo_type="dataset") | |
| n0 = len(examples) | |
| for r in map(json.loads, open(rp2, encoding="utf-8")): | |
| if r["status"] == "correct" and r["qid"] in split_rows: | |
| s = split_rows[r["qid"]] | |
| examples.append(TrainingExample( | |
| trajectory_id=r["qid"] + "-replay", | |
| original_failure_type="success_replay", root_cause="none", | |
| prompt=[{"role": "", "content": templ( | |
| PROMPT.format(schema=s["schema"], question=s["question"]))}], | |
| chosen=[{"role": "", "content": _extract(r["raw_output"])}], | |
| )) | |
| print(f"[dpo] +{len(examples)-n0} success-replay rows (label=True)", flush=True) | |
| cfg = RetrainConfig( | |
| model=args.model, | |
| algorithm=args.algorithm, | |
| output_dir=os.path.join(args.out, "adapter"), | |
| lora_r=args.lora_r, | |
| lora_alpha=args.lora_r * 2, | |
| lora_dropout=0.05, | |
| lora_target_modules=["q_proj", "k_proj", "v_proj", "o_proj", | |
| "gate_proj", "up_proj", "down_proj"], | |
| num_train_epochs=args.epochs, | |
| per_device_train_batch_size=args.batch_size, | |
| gradient_accumulation_steps=args.grad_accum, | |
| learning_rate=args.lr, | |
| beta=args.beta, | |
| max_steps=args.max_steps, | |
| extra={ | |
| "bf16": True, | |
| "logging_steps": 2, | |
| "save_strategy": "no", | |
| "report_to": "none", | |
| "gradient_checkpointing": True, | |
| }, | |
| ) | |
| if args.algorithm == "dpo": | |
| # The BCO wrapper instantiates the model itself and rejects model_init_kwargs. | |
| cfg.extra["model_init_kwargs"] = {"dtype": "bfloat16"} | |
| # Cohere vocab is ~256k: full-sequence fp32 logits are ~4 GiB per copy and OOM a | |
| # 24 GB card. use_logits_to_keep restricts logits to completion tokens (~60 here). | |
| # If this TRL lacks it, shrink sequence caps hard instead. | |
| import inspect | |
| import trl | |
| cfg_keys = set(inspect.signature(trl.DPOConfig.__init__).parameters) | |
| if "use_logits_to_keep" in cfg_keys: | |
| cfg.extra.update({"use_logits_to_keep": True, | |
| "max_prompt_length": 2944, "max_length": 3072}) | |
| print("[dpo] use_logits_to_keep=True (trl", trl.__version__, ")", flush=True) | |
| else: | |
| cfg.extra.update({"max_prompt_length": args.max_prompt_len, | |
| "max_length": args.max_prompt_len + 128}) | |
| print("[dpo] trl lacks use_logits_to_keep -> caps", args.max_prompt_len, flush=True) | |
| # rpo_alpha adds an NLL term on the chosen completion (DPO+SFT hybrid): the model | |
| # must LEARN to produce the verified fix, not merely prefer it. v1 (pure DPO, | |
| # 30 steps) failed to move pool accuracy at all - this is the main v2 fix. | |
| if args.rpo_alpha and "rpo_alpha" in cfg_keys: | |
| cfg.extra["rpo_alpha"] = args.rpo_alpha | |
| print("[dpo] rpo_alpha =", args.rpo_alpha, flush=True) | |
| elif args.rpo_alpha: | |
| print("[dpo] WARNING: trl lacks rpo_alpha; pure DPO", flush=True) | |
| lens = sorted(len(tok(templ(p["prompt_user_content"]))["input_ids"]) for p in healed) | |
| print(f"[dpo] prompt tokens: median={lens[len(lens)//2]} p90={lens[int(len(lens)*.9)]} " | |
| f"max={lens[-1]} | cap={cfg.extra['max_prompt_length']}", flush=True) | |
| trainer = build_retrainer(examples, cfg) | |
| stats = trainer.train() | |
| print("[dpo] train() returned:", str(stats)[:500], flush=True) | |
| # Save the adapter explicitly (find the underlying HF trainer/model). | |
| os.makedirs(cfg.output_dir, exist_ok=True) | |
| inner = None | |
| for attr in ("trainer", "_trainer", "dpo_trainer"): | |
| inner = getattr(trainer, attr, None) | |
| if inner is not None and hasattr(inner, "model"): | |
| break | |
| model_obj = getattr(inner, "model", None) or getattr(trainer, "model", None) | |
| if model_obj is None: | |
| raise RuntimeError( | |
| f"cannot locate trained model on {type(trainer).__name__}; attrs: " | |
| + ", ".join(a for a in dir(trainer) if not a.startswith("__"))[:800] | |
| ) | |
| model_obj.save_pretrained(cfg.output_dir) # PeftModel -> adapter only | |
| print("[dpo] adapter saved:", os.listdir(cfg.output_dir), flush=True) | |
| with open(os.path.join(args.out, "train_stats.json"), "w", encoding="utf-8") as f: | |
| json.dump({"n_pairs": len(healed), "stats": str(stats)[:2000], | |
| "config": {k: v for k, v in vars(args).items()}}, f, indent=2) | |
| api = HfApi() | |
| if not args.no_push_adapter: | |
| api.create_repo(args.adapter_repo, private=True, exist_ok=True) | |
| api.upload_folder(repo_id=args.adapter_repo, folder_path=cfg.output_dir, | |
| commit_message=f"DPO LoRA ({len(healed)} self-heal pairs)") | |
| print("[dpo] adapter pushed:", args.adapter_repo, flush=True) | |
| if args.push_repo: | |
| api.upload_folder(repo_id=args.push_repo, repo_type="dataset", | |
| folder_path=args.out, path_in_repo=args.push_path or "dpo", | |
| commit_message="dpo training stats") | |
| if __name__ == "__main__": | |
| main() | |