tools/train_schema_v2_synthetic.py

# -*- 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()