"""Bounded German multi-turn RL on top of a local SFT checkpoint.""" from __future__ import annotations import argparse from datetime import datetime, timezone import json import os from pathlib import Path import subprocess import sys from .agent_rl_tasks import generate_tasks from .evaluation import checkpoint_sha256, compare as compare_math, sha256_file def write_json(path: Path, value: dict) -> None: path.write_text(json.dumps(value, indent=2, ensure_ascii=False, allow_nan=False) + "\n") def source_hashes() -> dict: return {name: sha256_file(Path(__file__).with_name(name + ".py")) for name in ( "agent_rl_tasks", "agent_rl_training", "agent_rl_evaluation", "agent_rl_pilot", )} def prepare(model: str | Path, output: str | Path, *, retention_run: str | Path, steps: int = 32, train_size: int = 128, dev_size: int = 32, test_size: int = 32, seed: int = 42) -> dict: """Freeze synthetic split identities and all protocol inputs before evaluation.""" for name, value in (("steps", steps), ("train_size", train_size), ("dev_size", dev_size), ("test_size", test_size)): if type(value) is not int or value < 1: raise ValueError(f"{name} must be a positive integer") if type(seed) is not int or seed < 0: raise ValueError("seed must be a nonnegative integer") model, output, retention = (Path(p).resolve() for p in (model, output, retention_run)) if output.exists(): raise ValueError("Choose a fresh output directory") if not (model / "config.json").is_file() or (model / "adapter_config.json").exists(): raise ValueError("Supply a standalone local SFT checkpoint") for name in ("candidate/report.json", "candidate-spec.json", "evaluation.jsonl"): if not (retention / name).is_file(): raise ValueError(f"Missing prior mathematical retention evidence: {name}") from .evaluation import _verify_report prior = json.loads((retention / "candidate/report.json").read_text()) _verify_report(prior) digest = checkpoint_sha256(model) if prior["provenance"]["spec"]["adapter"]["sha256"] != digest: raise ValueError("Retention baseline must evaluate this exact starting checkpoint") splits = {split: generate_tasks(split, count, seed=seed) for split, count in ( ("train", train_size), ("dev", dev_size), ("test", test_size), )} ids = [row["id"] for rows in splits.values() for row in rows] if len(ids) != len(set(ids)): raise ValueError("Task identities must be disjoint across splits") from transformers import AutoTokenizer tokenizer = AutoTokenizer.from_pretrained(str(model), local_files_only=True, trust_remote_code=False) sources = source_hashes() output.mkdir(parents=True) tokenizer_dir = output / "frozen-tokenizer" tokenizer.save_pretrained(tokenizer_dir) for split, rows in splits.items(): (output / f"{split}.jsonl").write_text("".join( json.dumps(row, ensure_ascii=False, allow_nan=False) + "\n" for row in rows )) training = { "schema_version": 1, "model": {"path": str(model), "sha256": digest}, "dataset": {"path": str(output / "train.jsonl"), "sha256": sha256_file(output / "train.jsonl")}, "output_dir": str(output / "training"), "seed": seed, "chat_template_kwargs": {"enable_thinking": True}, "training": {"max_steps": steps, "num_generations": 4, "per_device_train_batch_size": 4, "gradient_accumulation_steps": 1, "max_completion_length": 768, "max_tool_calling_iterations": 3, "learning_rate": 1e-5, "beta": 0.02, "temperature": 0.6, "top_p": 0.95, "top_k": 20, "min_p": 0.0}, "lora": {"r": 16, "alpha": 32, "dropout": 0.0, "target_modules": "all-linear"}, } write_json(output / "training.json", training) plan = { "schema_version": 1, "status": "prepared", "model": str(model), "model_sha256": digest, "tokenizer_sha256": checkpoint_sha256(tokenizer_dir), "training_config_sha256": sha256_file(output / "training.json"), "split_sha256": {split: sha256_file(output / f"{split}.jsonl") for split in splits}, "split_sizes": {split: len(rows) for split, rows in splits.items()}, "source_sha256": sources, "seed": seed, "steps": steps, "sampling": {"temperature": 0.6, "top_p": 0.95, "top_k": 20, "min_p": 0.0, "seed": seed}, "chat_template_kwargs": {"enable_thinking": True}, "max_completion_length": 768, "max_tool_calling_iterations": 3, "retention_run": str(retention), "retention_baseline_report_sha256": sha256_file(retention / "candidate/report.json"), "retention_baseline_spec_sha256": sha256_file(retention / "candidate-spec.json"), "retention_tasks_sha256": sha256_file(retention / "evaluation.jsonl"), "notice": "Synthetic German tool-workflow pilot, not an open benchmark or general agent capability evaluation.", "release_eligible": False, } write_json(output / "plan.json", plan) return plan def evaluation_spec(plan: dict, root: Path, model: Path) -> dict: return { "schema_version": 1, "model": {"path": str(model), "sha256": checkpoint_sha256(model)}, "tokenizer": {"path": str(root / "frozen-tokenizer"), "sha256": plan["tokenizer_sha256"]}, "tasks": {"path": str(root / "test.jsonl"), "sha256": plan["split_sha256"]["test"]}, "sampling": plan["sampling"], "trials": 1, "max_completion_length": plan["max_completion_length"], "max_tool_calling_iterations": plan["max_tool_calling_iterations"], "chat_template_kwargs": plan["chat_template_kwargs"], "device": "cuda", "dtype": "bfloat16", } def verify_inputs(root: Path, plan: dict) -> None: expected = { "model": (plan["model_sha256"], checkpoint_sha256(Path(plan["model"]))), "tokenizer": (plan["tokenizer_sha256"], checkpoint_sha256(root / "frozen-tokenizer")), "config": (plan["training_config_sha256"], sha256_file(root / "training.json")), "sources": (plan["source_sha256"], source_hashes()), } for split, digest in plan["split_sha256"].items(): expected[split] = (digest, sha256_file(root / f"{split}.jsonl")) retention = Path(plan["retention_run"]) for key, name in (("report", "candidate/report.json"), ("spec", "candidate-spec.json"), ("tasks", "evaluation.jsonl")): field = f"retention_baseline_{key}_sha256" if key != "tasks" else "retention_tasks_sha256" expected[f"retention_{key}"] = (plan[field], sha256_file(retention / name)) for name, (frozen, current) in expected.items(): if frozen != current: raise ValueError(f"Prepared input changed: {name}") def require_learning_update(report: dict) -> None: """A completed optimizer loop alone does not demonstrate reinforcement learning.""" evidence = report.get("update_evidence", {}) if report.get("status") != "completed" or any( type(evidence.get(name)) is not int or evidence[name] < 1 for name in ( "changed_trainable_tensors", "mixed_reward_groups", "nonzero_gradient_steps", ) ): raise RuntimeError("No verified reward-driven parameter update; inspect training-only rollouts") def execute(plan_path: str | Path) -> dict: plan_path = Path(plan_path).resolve() root = plan_path.parent plan = json.loads(plan_path.read_text()) status_path = root / "run-manifest.json" if status_path.exists(): raise ValueError("Run already started; choose a fresh prepared output directory") import torch if not torch.cuda.is_available() or not torch.cuda.is_bf16_supported(): raise RuntimeError("Run this pilot on a CUDA GPU allocation with BF16 support") verify_inputs(root, plan) logs = root / "logs" logs.mkdir() status = {"status": "running", "started_at": datetime.now(timezone.utc).isoformat(), "plan_sha256": sha256_file(plan_path), "completed_stages": [], "gpu": torch.cuda.get_device_name(), "release_eligible": False} env = {**os.environ, "HF_HUB_OFFLINE": "1", "TRANSFORMERS_OFFLINE": "1", "TOKENIZERS_PARALLELISM": "false", "OMP_NUM_THREADS": "4"} def stage(name: str, arguments: list[str]) -> None: verify_inputs(root, plan) status["stage"] = name write_json(status_path, status) with (logs / f"{name}.log").open("w") as log: subprocess.run([sys.executable, *arguments], env=env, stdout=log, stderr=subprocess.STDOUT, check=True) status["completed_stages"].append(name) write_json(status_path, status) try: base, candidate = Path(plan["model"]), root / "merged-checkpoint" write_json(root / "baseline-spec.json", evaluation_spec(plan, root, base)) stage("baseline", ["-m", "hessian_recipes.agent_rl_evaluation", str(root / "baseline-spec.json"), "--output", str(root / "baseline")]) from .agent_rl_evaluation import compare, verify_report verify_report(json.loads((root / "baseline/report.json").read_text())) stage("grpo", ["-m", "hessian_recipes.agent_rl_training", str(root / "training.json")]) require_learning_update(json.loads((root / "training/training-manifest.json").read_text())) stage("merge", ["-m", "hessian_recipes.checkpoints", str(base), str(root / "training/checkpoint"), str(candidate), "--dtype", "bfloat16", "--device", "cuda"]) write_json(root / "candidate-spec.json", evaluation_spec(plan, root, candidate)) stage("candidate", ["-m", "hessian_recipes.agent_rl_evaluation", str(root / "candidate-spec.json"), "--output", str(root / "candidate")]) verify_report(json.loads((root / "candidate/report.json").read_text())) comparison = compare(root / "baseline/report.json", root / "candidate/report.json") if comparison.get("paired") is None: raise ValueError("Agent reports cannot form a verified matched comparison") write_json(root / "comparison.json", comparison) retention = Path(plan["retention_run"]) spec = json.loads((retention / "candidate-spec.json").read_text()) digest = checkpoint_sha256(candidate) spec["benchmark"]["path"] = str(retention / "evaluation.jsonl") spec["model"].update(id="local-qwen3-1.7b-german-agent-rl", revision=digest) spec["adapter"].update(path=str(candidate), sha256=digest) write_json(root / "retention-spec.json", spec) stage("retention", ["-m", "hessian_recipes", "eval", str(root / "retention-spec.json"), "--output", str(root / "retention")]) from .evaluation import _verify_report _verify_report(json.loads((root / "retention/report.json").read_text())) math_comparison = compare_math(retention / "candidate/report.json", root / "retention/report.json") write_json(root / "retention-comparison.json", math_comparison) status.update(status="completed", comparison=comparison, retention_comparison=math_comparison, completed_at=datetime.now(timezone.utc).isoformat()) except BaseException as error: status.update(status="failed", error={"type": type(error).__name__, "message": str(error)}) raise finally: write_json(status_path, status) return status def main() -> None: parser = argparse.ArgumentParser(description=__doc__) commands = parser.add_subparsers(dest="command", required=True) prep = commands.add_parser("prepare") prep.add_argument("--model", type=Path, required=True) prep.add_argument("--output", type=Path, required=True) prep.add_argument("--retention-run", type=Path, required=True) prep.add_argument("--steps", type=int, default=32) prep.add_argument("--train-size", type=int, default=128) prep.add_argument("--dev-size", type=int, default=32) prep.add_argument("--test-size", type=int, default=32) prep.add_argument("--seed", type=int, default=42) run = commands.add_parser("run") run.add_argument("--plan", type=Path, required=True) args = parser.parse_args() result = (prepare(args.model, args.output, retention_run=args.retention_run, steps=args.steps, train_size=args.train_size, dev_size=args.dev_size, test_size=args.test_size, seed=args.seed) if args.command == "prepare" else execute(args.plan)) print(json.dumps(result, indent=2, ensure_ascii=False)) if __name__ == "__main__": main()