"""Bounded German Qwen SFT pilot, with frozen data and matched GPU evaluations.""" from __future__ import annotations import argparse from datetime import datetime, timezone import json import os from pathlib import Path import random import subprocess import sys from .evaluation import checkpoint_sha256, compare, sha256_file from .training import read_jsonl_snapshot def write_json(path: Path, value: dict) -> None: path.write_text(json.dumps(value, indent=2, ensure_ascii=False, allow_nan=False) + "\n") def write_rows(path: Path, rows: list[dict]) -> None: path.write_text("".join(json.dumps(row, ensure_ascii=False) + "\n" for row in rows)) def validate_splits(train: list[dict], dev: list[dict], evaluation: list[dict]) -> None: seen = [] for name, rows in (("train", train), ("dev", dev), ("evaluation", evaluation)): if not rows: raise ValueError(f"{name} split is empty") ids = [row["source_id"] for row in rows] if len(ids) != len(set(ids)): raise ValueError(f"{name} has duplicate source identities") seen.append(set(ids)) if any(seen[a] & seen[b] for a, b in ((0, 1), (0, 2), (1, 2))): raise ValueError("Training, development and evaluation source identities must be disjoint") def prepare(model: str | Path, data: str | Path, output: str | Path, *, steps: int = 128, eval_size: int = 64, max_length: int = 1024, max_new_tokens: int = 768, seed: int = 42) -> dict: for name, value in (("steps", steps), ("eval_size", eval_size), ("max_length", max_length), ("max_new_tokens", max_new_tokens)): if type(value) is not int or value < 1: raise ValueError(f"{name} must be a positive integer") model, data, output = Path(model).resolve(), Path(data).resolve(), Path(output).resolve() if not model.is_dir(): raise ValueError("Supply an existing local model checkpoint") if output.exists(): raise ValueError("Choose a fresh output directory") train, train_sha = read_jsonl_snapshot(data / "train.jsonl") dev, dev_sha = read_jsonl_snapshot(data / "dev.jsonl") external, external_sha = read_jsonl_snapshot(data / "mgsm-de.jsonl") validate_splits(train, dev, external) if eval_size > len(external): raise ValueError("eval_size exceeds the available external evaluation set") from transformers import AutoTokenizer tokenizer = AutoTokenizer.from_pretrained(str(model), local_files_only=True, trust_remote_code=False) mode = {"enable_thinking": True} retained, rejected = [], [] for row in train: tokens = tokenizer.apply_chat_template( [{"role": "user", "content": row["prompt"]}, {"role": "assistant", "content": row["completion"]}], tokenize=True, return_dict=False, add_generation_prompt=False, **mode, ) if len(tokens) <= max_length: retained.append(row) else: rejected.append({"source_id": row["source_id"], "tokens": len(tokens)}) if not retained: raise ValueError("All training examples exceed the token budget") chosen = random.Random(seed).sample(external, eval_size) output.mkdir(parents=True) tokenizer_dir = output / "frozen-tokenizer" tokenizer.save_pretrained(tokenizer_dir) write_rows(output / "train.jsonl", retained) write_rows(output / "dev.jsonl", dev) write_rows(output / "evaluation.jsonl", chosen) recipe = { "schema_version": 1, "method": "sft", "model": {"name": str(model)}, "dataset": "train.jsonl", "output_dir": "training", "seed": seed, "chat_template_kwargs": mode, "lora": {"r": 16, "alpha": 32, "dropout": 0.0, "target_modules": "all-linear"}, "training": {"max_steps": steps, "max_length": max_length, "per_device_train_batch_size": 1, "gradient_accumulation_steps": 8, "learning_rate": 1e-4, "bf16": True, "use_cpu": False, "gradient_checkpointing": True, "logging_steps": 4}, } write_json(output / "sft.json", recipe) plan = { "schema_version": 1, "status": "prepared", "model": str(model), "model_sha256": checkpoint_sha256(model), "seed": seed, "recipe_sha256": sha256_file(output / "sft.json"), "train_sha256": sha256_file(output / "train.jsonl"), "evaluation_sha256": sha256_file(output / "evaluation.jsonl"), "tokenizer_sha256": checkpoint_sha256(tokenizer_dir), "source_data": {"train_sha256": train_sha, "dev_sha256": dev_sha, "mgsm_sha256": external_sha, "manifest_sha256": sha256_file(data / "manifest.json")}, "train_examples": len(retained), "overlength_rejected": rejected, "development_examples": len(dev), "external_evaluation_examples": len(chosen), "external_evaluation_total": len(external), "steps": steps, "max_new_tokens": max_new_tokens, "chat_template": tokenizer.chat_template, "chat_template_kwargs": mode, "notice": "German mathematical-reasoning pilot. Source grouping prevents split leakage; translation quality and general German reasoning are not established by numerical accuracy.", } write_json(output / "plan.json", plan) return plan def evaluation_spec(plan: dict, root: Path, checkpoint: Path, label: str) -> dict: digest = checkpoint_sha256(checkpoint) return { "schema_version": 1, "benchmark": {"id": f"mgsm-de-pilot-{plan['external_evaluation_examples']}", "version": "2026-10-15-v1", "status": "research-pilot", "count": plan["external_evaluation_examples"], "path": "evaluation.jsonl", "sha256": plan["evaluation_sha256"]}, "model": {"id": f"local-qwen3-1.7b-{label}", "revision": digest, "tokenizer_id": "frozen-qwen3-tokenizer", "tokenizer_revision": plan["tokenizer_sha256"], "chat_template": plan["chat_template"], "chat_template_kwargs": plan["chat_template_kwargs"]}, "sampling": {"temperature": 0.6, "top_p": 0.95, "top_k": 20, "min_p": 0.0, "max_new_tokens": plan["max_new_tokens"], "seed": plan["seed"]}, "trials": 1, "scorer_version": "1", "harness_version": "1", "adapter": {"type": "transformers", "path": str(checkpoint), "sha256": digest, "tokenizer_path": str(root / "frozen-tokenizer"), "tokenizer_sha256": plan["tokenizer_sha256"], "device": "cuda", "dtype": "bfloat16"}, } def execute(plan_path: str | Path) -> dict: """Every model-consuming stage runs in a separate process on the GPU worker.""" 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; use 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") for name, actual in ( ("model_sha256", checkpoint_sha256(Path(plan["model"]))), ("recipe_sha256", sha256_file(root / "sft.json")), ("train_sha256", sha256_file(root / "train.jsonl")), ("evaluation_sha256", sha256_file(root / "evaluation.jsonl")), ("tokenizer_sha256", checkpoint_sha256(root / "frozen-tokenizer")), ): if plan[name] != actual: raise ValueError(f"Prepared input changed: {name}") logs = root / "logs" logs.mkdir() status = {"status": "running", "started_at": datetime.now(timezone.utc).isoformat(), "plan_sha256": sha256_file(plan_path), "gpu": torch.cuda.get_device_name(), "completed_stages": [], "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: 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: write_json(root / "baseline-spec.json", evaluation_spec(plan, root, Path(plan["model"]), "baseline")) stage("baseline", ["-m", "hessian_recipes", "eval", str(root / "baseline-spec.json"), "--output", str(root / "baseline")]) stage("sft", ["-m", "hessian_recipes", "train", str(root / "sft.json")]) stage("merge", ["-m", "hessian_recipes.checkpoints", plan["model"], str(root / "training/checkpoint"), str(root / "merged-checkpoint"), "--dtype", "bfloat16", "--device", "cuda"]) write_json(root / "candidate-spec.json", evaluation_spec(plan, root, root / "merged-checkpoint", "adapted")) stage("candidate", ["-m", "hessian_recipes", "eval", str(root / "candidate-spec.json"), "--output", str(root / "candidate")]) gate = compare(root / "baseline/report.json", root / "candidate/report.json") write_json(root / "comparison.json", gate) status.update(status="completed", comparison=gate, 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() -> int: 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("--data", type=Path, required=True) prep.add_argument("--output", type=Path, required=True) prep.add_argument("--steps", type=int, default=128) prep.add_argument("--eval-size", type=int, default=64) prep.add_argument("--max-length", type=int, default=1024) prep.add_argument("--max-new-tokens", type=int, default=768) run = commands.add_parser("run") run.add_argument("--plan", type=Path, required=True) args = parser.parse_args() result = prepare(args.model, args.data, args.output, steps=args.steps, eval_size=args.eval_size, max_length=args.max_length, max_new_tokens=args.max_new_tokens) if args.command == "prepare" else execute(args.plan) print(json.dumps(result, indent=2)) return 0 if __name__ == "__main__": raise SystemExit(main())