Add Rust cached synthetic training runner

AGENTS.md

@@ -117,29 +117,31 @@ cargo run --release --manifest-path tools\schema_v2_synthetic_augment\Cargo.toml
 ```
 
 Preferred synthetic follow-up training is a second stage from the best repaired
-hard-focus checkpoint, not a replacement for hard-focus. Do not combine
-`--encoded-cache-dir` with `--extra-data-file`; use the raw hard-focus JSONL
-when mixing synthetic augmentation, or rebuild a combined Rust encoded cache.
-Use native Windows Python from `.venv` after confirming CUDA works:
+hard-focus checkpoint, not a replacement for hard-focus. Keep this path
+Rust-cache-first: build one combined encoded cache from hard-focus JSONL plus
+synthetic JSONL, then train from that cache. Do not pass `--extra-data-file` to
+`anifilebert.train` together with `--encoded-cache-dir`.
+
+Use the local wrapper, which calls Rust `tools/encoded_dataset_cache` with
+multiple `--input` values and then launches `anifilebert.train` against the
+combined cache:
 
 ```powershell
-.\.venv\Scripts\python.exe -m anifilebert.train --tokenizer char `
-  --data-file data\schema_v2_hard_focus_char_seed63.jsonl `
-  --extra-data-file data\schema_v2_synthetic_aug.jsonl `
-  --extra-data-repeat 3 `
-  --vocab-file datasets\AnimeName\vocab.char.json `
-  --save-dir checkpoints\schema-v2-best-hardfocus-synth-pathleaf `
-  --init-model-dir checkpoints\ablation-schema-v2-hardfocus-cache-repaired-from-baseline-seed62-10epoch-rerun\final `
-  --epochs 2 --batch-size 512 --learning-rate 0.00004 --warmup-steps 120 `
-  --max-seq-length 128 --train-split 0.995 --num-workers 0 `
-  --checkpoint-steps 1000 --save-total-limit 3 --no-periodic-eval `
-  --bf16 --auto-find-batch-size `
-  --parse-eval-limit 2048 `
-  --case-eval-file data\parser_regression_cases.json `
-  --case-eval-output reports\schema_v2_best_hardfocus_synth_pathleaf_case_metrics.json `
-  --seed 63 --experiment-name schema-v2-best-hardfocus-synth-pathleaf
+.\.venv\Scripts\python.exe -m tools.train_schema_v2_synthetic
 ```
 
+The wrapper defaults to:
+
+- primary data: `data\schema_v2_hard_focus_char_seed63.jsonl`
+- synthetic data: `data\schema_v2_synthetic_aug.jsonl`
+- synthetic repeat: `3`
+- encoded cache: `data\encoded_cache\schema_v2_hard_focus_seed63_synth_pathleaf_repeat3`
+- init checkpoint: `checkpoints\ablation-schema-v2-hardfocus-cache-repaired-from-baseline-seed62-10epoch-rerun\final`
+- output checkpoint: `checkpoints\schema-v2-best-hardfocus-synth-pathleaf-cache`
+
+Use `--force-cache` to rebuild the combined cache after changing either JSONL,
+vocab, label schema, max length, split ratio, seed, or repeat count.
+
 Export for Android:
 
 ```bash

tools/encoded_dataset_cache/README.md

@@ -21,6 +21,26 @@ cargo run --release --manifest-path tools\encoded_dataset_cache\Cargo.toml -- `
   --threads 16
 ```
 
+Multiple JSONL inputs can be encoded into one deterministic train/eval split.
+Pass `--input-repeat` once per `--input` when an augmentation source should be
+upweighted:
+
+```powershell
+cargo run --release --manifest-path tools\encoded_dataset_cache\Cargo.toml -- `
+  --input data\schema_v2_hard_focus_char_seed63.jsonl `
+  --input data\schema_v2_synthetic_aug.jsonl `
+  --input-repeat 1 `
+  --input-repeat 3 `
+  --vocab-file datasets\AnimeName\vocab.char.json `
+  --label-schema-file label_schema.json `
+  --output-dir data\encoded_cache\schema_v2_hard_focus_seed63_synth_pathleaf_repeat3 `
+  --max-length 128 `
+  --train-split 0.995 `
+  --seed 63 `
+  --shard-size 25000 `
+  --threads 16
+```
+
 Use the cache in training:
 
 ```powershell

tools/encoded_dataset_cache/src/main.rs

@@ -1,3 +1,5 @@
+mod regex_benchmark;
+
 use anyhow::{bail, Context, Result};
 use clap::Parser;
 use fancy_regex::Regex as FancyRegex;
@@ -78,13 +80,16 @@ const SEPARATOR_CHARS: &[char] = &[' ', '\t', '-', '_', '.', '|', '~', '～'];
 )]
 struct Args {
     #[arg(long)]
-    input: PathBuf,
+    input: Vec<PathBuf>,
+
+    #[arg(long, value_name = "N")]
+    input_repeat: Vec<usize>,
 
     #[arg(long)]
-    vocab_file: PathBuf,
+    vocab_file: Option<PathBuf>,
 
     #[arg(long)]
-    output_dir: PathBuf,
+    output_dir: Option<PathBuf>,
 
     #[arg(long, default_value = "label_schema.json")]
     label_schema_file: PathBuf,
@@ -109,6 +114,15 @@ struct Args {
 
     #[arg(long, default_value_t = 0)]
     threads: usize,
+
+    #[arg(long)]
+    regex_benchmark_input: Option<PathBuf>,
+
+    #[arg(long, default_value_t = 0)]
+    regex_benchmark_limit_rows: usize,
+
+    #[arg(long, default_value_t = 3)]
+    regex_benchmark_repeat: usize,
 }
 
 #[derive(Debug, Deserialize)]
@@ -160,6 +174,24 @@ struct SplitSummary {
 
 fn main() -> Result<()> {
     let args = Args::parse();
+    if let Some(input) = &args.regex_benchmark_input {
+        return regex_benchmark::run(
+            input,
+            args.regex_benchmark_limit_rows,
+            args.regex_benchmark_repeat,
+        );
+    }
+    if args.input.is_empty() {
+        bail!("at least one --input is required");
+    }
+    let vocab_file = args
+        .vocab_file
+        .as_ref()
+        .context("--vocab-file is required when building an encoded cache")?;
+    let output_dir = args
+        .output_dir
+        .as_ref()
+        .context("--output-dir is required when building an encoded cache")?;
     if args.max_length < 4 {
         bail!("--max-length must be at least 4");
     }
@@ -177,9 +209,10 @@ fn main() -> Result<()> {
     }
 
     let started = Instant::now();
-    let vocab = load_vocab(&args.vocab_file)?;
+    let vocab = load_vocab(vocab_file)?;
     let label_ids = load_label_ids(&args.label_schema_file)?;
-    let mut rows = load_rows(&args.input, args.limit_rows)?;
+    let input_repeats = resolve_input_repeats(&args.input, &args.input_repeat)?;
+    let (mut rows, input_summaries) = load_input_rows(&args.input, &input_repeats, args.limit_rows)?;
     if rows.len() < 2 {
         bail!("need at least two rows to build train/eval cache");
     }
@@ -192,10 +225,10 @@ fn main() -> Result<()> {
     let split_idx = split_idx.max(1).min(rows.len() - 1);
     let (train_rows, eval_rows) = rows.split_at(split_idx);
 
-    fs::create_dir_all(&args.output_dir).with_context(|| {
+    fs::create_dir_all(output_dir).with_context(|| {
         format!(
             "failed to create output directory {}",
-            args.output_dir.display()
+            output_dir.display()
         )
     })?;
 
@@ -207,25 +240,26 @@ fn main() -> Result<()> {
     let train_summary = write_split(
         "train",
         train_rows,
-        &args.output_dir,
+        output_dir,
         &context,
         args.shard_size,
     )?;
     let eval_summary = write_split(
         "eval",
         eval_rows,
-        &args.output_dir,
+        output_dir,
         &context,
         args.shard_size,
     )?;
-    write_eval_records(eval_rows, &args.output_dir.join("eval_records.jsonl"))?;
+    write_eval_records(eval_rows, &output_dir.join("eval_records.jsonl"))?;
 
     let manifest = json!({
         "format": "anifilebert.encoded_dataset_cache.v1",
-        "input": args.input,
-        "vocab_file": args.vocab_file,
+        "input": args.input.first(),
+        "inputs": input_summaries,
+        "vocab_file": vocab_file,
         "label_schema_file": args.label_schema_file,
-        "output_dir": args.output_dir,
+        "output_dir": output_dir,
         "max_length": args.max_length,
         "shard_size": args.shard_size,
         "limit_rows": args.limit_rows,
@@ -238,7 +272,7 @@ fn main() -> Result<()> {
         "eval_records": "eval_records.jsonl",
         "elapsed_seconds": started.elapsed().as_secs_f64(),
     });
-    let manifest_path = args.output_dir.join("manifest.json");
+    let manifest_path = output_dir.join("manifest.json");
     fs::write(&manifest_path, serde_json::to_string_pretty(&manifest)?)
         .with_context(|| format!("failed to write {}", manifest_path.display()))?;
     println!("{}", serde_json::to_string_pretty(&manifest)?);
@@ -293,14 +327,67 @@ fn load_label_ids(path: &Path) -> Result<HashMap<String, i16>> {
         .collect())
 }
 
-fn load_rows(path: &Path, limit_rows: usize) -> Result<Vec<SourceRow>> {
+fn resolve_input_repeats(inputs: &[PathBuf], repeats: &[usize]) -> Result<Vec<usize>> {
+    if repeats.is_empty() {
+        return Ok(vec![1; inputs.len()]);
+    }
+    if repeats.len() == 1 {
+        return Ok(vec![repeats[0].max(1); inputs.len()]);
+    }
+    if repeats.len() != inputs.len() {
+        bail!(
+            "--input-repeat must be omitted, passed once for all inputs, or passed once per --input ({} inputs, {} repeats)",
+            inputs.len(),
+            repeats.len()
+        );
+    }
+    Ok(repeats.iter().map(|repeat| (*repeat).max(1)).collect())
+}
+
+fn load_input_rows(
+    inputs: &[PathBuf],
+    repeats: &[usize],
+    limit_rows: usize,
+) -> Result<(Vec<SourceRow>, Vec<Value>)> {
+    let mut combined = Vec::new();
+    let mut summaries = Vec::new();
+    for (path, repeat) in inputs.iter().zip(repeats.iter()) {
+        let rows = load_rows(path)?;
+        let samples = rows.len();
+        let mut written = 0usize;
+        for _ in 0..*repeat {
+            for row in &rows {
+                if limit_rows > 0 && combined.len() >= limit_rows {
+                    break;
+                }
+                let mut row = row.clone();
+                row.row_index = combined.len();
+                combined.push(row);
+                written += 1;
+            }
+            if limit_rows > 0 && combined.len() >= limit_rows {
+                break;
+            }
+        }
+        summaries.push(json!({
+            "path": path,
+            "samples": samples,
+            "repeat": repeat,
+            "effective_samples": samples * repeat,
+            "written_rows": written,
+        }));
+        if limit_rows > 0 && combined.len() >= limit_rows {
+            break;
+        }
+    }
+    Ok((combined, summaries))
+}
+
+fn load_rows(path: &Path) -> Result<Vec<SourceRow>> {
     let file = File::open(path).with_context(|| format!("failed to open {}", path.display()))?;
     let reader = BufReader::new(file);
     let mut rows = Vec::new();
     for (idx, line) in reader.lines().enumerate() {
-        if limit_rows > 0 && rows.len() >= limit_rows {
-            break;
-        }
         let raw_line = line.with_context(|| format!("failed reading line {}", idx + 1))?;
         if raw_line.trim().is_empty() {
             continue;

tools/encoded_dataset_cache/src/{bin/regex_benchmark.rs → regex_benchmark.rs}

@@ -1,5 +1,4 @@
 use anyhow::{ensure, Context, Result};
-use clap::Parser;
 use fancy_regex::Regex as FancyRegex;
 use regex::Regex;
 use serde_json::Value;
@@ -24,37 +23,21 @@ const CJK_MARKER_PATTERN: &str = r"(?:[一二三四五六七八九十兩两貳
 const SPECIAL_CONTEXT_PREFIX_PATTERN: &str =
     r"(?i)^(?:[\[\(【《]\s*(?:menu|menus|bdmenu|ncop|nced|op|ed|ova|oad|sp)\s*[\]\)】》]\s*){0,2}";
 
-#[derive(Parser, Debug)]
-#[command(
-    about = "Compare regex vs fancy-regex workload costs for AniFileBERT cache preprocessing"
-)]
-struct Args {
-    #[arg(long)]
-    input: PathBuf,
+pub fn run(input: &PathBuf, limit_rows: usize, repeat: usize) -> Result<()> {
+    ensure!(repeat > 0, "--regex-benchmark-repeat must be greater than 0");
 
-    #[arg(long, default_value_t = 0)]
-    limit_rows: usize,
-
-    #[arg(long, default_value_t = 3)]
-    repeat: usize,
-}
-
-fn main() -> Result<()> {
-    let args = Args::parse();
-    ensure!(args.repeat > 0, "--repeat must be greater than 0");
-
-    let filenames = load_filenames(&args.input, args.limit_rows)?;
+    let filenames = load_filenames(input, limit_rows)?;
     if filenames.is_empty() {
-        anyhow::bail!("no filenames loaded from {}", args.input.display());
+        anyhow::bail!("no filenames loaded from {}", input.display());
     }
 
     let selective = SelectivePatterns::new()?;
     let fancy_all = FancyAllPatterns::new()?;
 
     let (selective_seconds, selective_count) =
-        time_repeated(args.repeat, || run_selective(&filenames, &selective))?;
+        time_repeated(repeat, || run_selective(&filenames, &selective))?;
     let (fancy_seconds, fancy_count) =
-        time_repeated(args.repeat, || run_fancy_all(&filenames, &fancy_all))?;
+        time_repeated(repeat, || run_fancy_all(&filenames, &fancy_all))?;
     ensure!(
         selective_count == fancy_count,
         "selective and fancy-all match counts differ: selective={}, fancy_all={}",
@@ -71,7 +54,7 @@ fn main() -> Result<()> {
         "{}",
         serde_json::json!({
             "rows": filenames.len(),
-            "repeat": args.repeat,
+            "repeat": repeat,
             "selective_seconds": selective_seconds,
             "fancy_all_seconds": fancy_seconds,
             "ratio": ratio,

tools/train_schema_v2_synthetic.py

@@ -0,0 +1,180 @@
+# -*- 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
+
+
+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()