# -*- coding: utf-8 -*- r"""Local schema v2 synthetic-augmentation training runner. This wrapper keeps the training structure reproducible: 1. Build a combined Rust encoded cache from the hard-focus JSONL plus synthetic augmentation JSONL. 2. Train with ``anifilebert.train --encoded-cache-dir`` so Python training never has to re-split raw mixed JSONL in a non-comparable way. Typical usage from the repo root on the local Windows GPU machine: .\.venv\Scripts\python.exe -m tools.train_schema_v2_synthetic """ from __future__ import annotations import argparse import datetime as dt import json from pathlib import Path import shlex import shutil import subprocess import sys from typing import Any, Sequence try: sys.stdout.reconfigure(encoding="utf-8", errors="replace") sys.stderr.reconfigure(encoding="utf-8", errors="replace") except AttributeError: pass def utc_now() -> str: return dt.datetime.now(dt.timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z") def command_text(args: Sequence[Any]) -> str: return " ".join(shlex.quote(str(arg)) for arg in args) def run(args: Sequence[Any], *, dry_run: bool, command_log: list[dict[str, Any]]) -> None: entry: dict[str, Any] = { "cmd": command_text(args), "started_at": utc_now(), "dry_run": dry_run, } command_log.append(entry) print(f"\n$ {entry['cmd']}") if dry_run: entry["returncode"] = 0 entry["finished_at"] = utc_now() return proc = subprocess.Popen( list(map(str, args)), stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, encoding="utf-8", errors="replace", bufsize=1, ) assert proc.stdout is not None for line in proc.stdout: print(line, end="") proc.wait() entry["returncode"] = proc.returncode entry["finished_at"] = utc_now() if proc.returncode != 0: raise RuntimeError(f"Command failed with exit code {proc.returncode}: {entry['cmd']}") def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Train schema v2 hard-focus + synthetic augmentation with Rust cache") parser.add_argument("--primary-data-file", default="data/schema_v2_hard_focus_char_seed63.jsonl") parser.add_argument("--synthetic-data-file", default="data/schema_v2_synthetic_aug.jsonl") parser.add_argument("--synthetic-repeat", type=int, default=3) parser.add_argument("--vocab-file", default="datasets/AnimeName/vocab.char.json") parser.add_argument("--label-schema-file", default="label_schema.json") parser.add_argument("--encoded-cache-dir", default="data/encoded_cache/schema_v2_hard_focus_seed63_synth_pathleaf_repeat3") parser.add_argument("--save-dir", default="checkpoints/schema-v2-best-hardfocus-synth-pathleaf-cache") parser.add_argument("--init-model-dir", default="checkpoints/ablation-schema-v2-hardfocus-cache-repaired-from-baseline-seed62-10epoch-rerun/final") parser.add_argument("--case-eval-output", default="reports/schema_v2_best_hardfocus_synth_pathleaf_cache_case_metrics.json") parser.add_argument("--experiment-name", default="schema-v2-best-hardfocus-synth-pathleaf-cache") parser.add_argument("--max-length", type=int, default=128) parser.add_argument("--train-split", type=float, default=0.995) parser.add_argument("--seed", type=int, default=63) parser.add_argument("--shard-size", type=int, default=25000) parser.add_argument("--threads", type=int, default=16) parser.add_argument("--epochs", type=float, default=2) parser.add_argument("--batch-size", type=int, default=512) parser.add_argument("--learning-rate", type=float, default=0.00004) parser.add_argument("--warmup-steps", type=int, default=120) parser.add_argument("--checkpoint-steps", type=int, default=1000) parser.add_argument("--save-total-limit", type=int, default=3) parser.add_argument("--parse-eval-limit", type=int, default=2048) parser.add_argument("--case-eval-file", default="data/parser_regression_cases.json") parser.add_argument("--force-cache", action="store_true", help="Delete and rebuild the encoded cache even if manifest exists") parser.add_argument("--skip-cache", action="store_true", help="Reuse the existing encoded cache") parser.add_argument("--dry-run", action="store_true") return parser.parse_args() def main() -> None: args = parse_args() command_log: list[dict[str, Any]] = [] cache_dir = Path(args.encoded_cache_dir) manifest_path = cache_dir / "manifest.json" if args.force_cache and cache_dir.exists(): print(f"Removing existing cache: {cache_dir}") if not args.dry_run: shutil.rmtree(cache_dir) if not args.skip_cache and not manifest_path.exists(): cache_cmd = [ "cargo", "run", "--release", "--manifest-path", "tools/encoded_dataset_cache/Cargo.toml", "--", "--input", args.primary_data_file, "--input", args.synthetic_data_file, "--input-repeat", "1", "--input-repeat", str(max(1, args.synthetic_repeat)), "--vocab-file", args.vocab_file, "--label-schema-file", args.label_schema_file, "--output-dir", args.encoded_cache_dir, "--max-length", str(args.max_length), "--train-split", str(args.train_split), "--seed", str(args.seed), "--shard-size", str(args.shard_size), "--threads", str(args.threads), ] run(cache_cmd, dry_run=args.dry_run, command_log=command_log) else: print(f"Using existing encoded cache: {cache_dir}") train_cmd = [ sys.executable, "-m", "anifilebert.train", "--tokenizer", "char", "--data-file", args.primary_data_file, "--vocab-file", args.vocab_file, "--encoded-cache-dir", args.encoded_cache_dir, "--save-dir", args.save_dir, "--init-model-dir", args.init_model_dir, "--epochs", str(args.epochs), "--batch-size", str(args.batch_size), "--learning-rate", str(args.learning_rate), "--warmup-steps", str(args.warmup_steps), "--max-seq-length", str(args.max_length), "--train-split", str(args.train_split), "--num-workers", "0", "--checkpoint-steps", str(args.checkpoint_steps), "--save-total-limit", str(args.save_total_limit), "--no-periodic-eval", "--bf16", "--auto-find-batch-size", "--parse-eval-limit", str(args.parse_eval_limit), "--case-eval-file", args.case_eval_file, "--case-eval-output", args.case_eval_output, "--seed", str(args.seed), "--experiment-name", args.experiment_name, ] run(train_cmd, dry_run=args.dry_run, command_log=command_log) run_manifest = { "name": args.experiment_name, "started_at": command_log[0]["started_at"] if command_log else utc_now(), "finished_at": utc_now(), "primary_data_file": args.primary_data_file, "synthetic_data_file": args.synthetic_data_file, "synthetic_repeat": args.synthetic_repeat, "encoded_cache_dir": args.encoded_cache_dir, "save_dir": args.save_dir, "init_model_dir": args.init_model_dir, "commands": command_log, } manifest_output = Path(args.save_dir) / "schema_v2_synthetic_train_manifest.json" print(f"Writing run manifest: {manifest_output}") if not args.dry_run: manifest_output.parent.mkdir(parents=True, exist_ok=True) manifest_output.write_text(json.dumps(run_manifest, ensure_ascii=False, indent=2), encoding="utf-8") if __name__ == "__main__": main()