Add tests for combo generator

tools/virtual_dataset_generator/Cargo.lock

@@ -2,6 +2,15 @@
 # It is not intended for manual editing.
 version = 4
 
+[[package]]
+name = "aho-corasick"
+version = "1.1.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301"
+dependencies = [
+ "memchr",
+]
+
 [[package]]
 name = "anifilebert-virtual-dataset-generator"
 version = "0.1.0"
@@ -10,6 +19,7 @@ dependencies = [
  "clap",
  "rand",
  "rayon",
+ "regex",
  "serde",
  "serde_json",
 ]
@@ -277,6 +287,35 @@ dependencies = [
  "crossbeam-utils",
 ]
 
+[[package]]
+name = "regex"
+version = "1.12.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e10754a14b9137dd7b1e3e5b0493cc9171fdd105e0ab477f51b72e7f3ac0e276"
+dependencies = [
+ "aho-corasick",
+ "memchr",
+ "regex-automata",
+ "regex-syntax",
+]
+
+[[package]]
+name = "regex-automata"
+version = "0.4.14"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6e1dd4122fc1595e8162618945476892eefca7b88c52820e74af6262213cae8f"
+dependencies = [
+ "aho-corasick",
+ "memchr",
+ "regex-syntax",
+]
+
+[[package]]
+name = "regex-syntax"
+version = "0.8.10"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "dc897dd8d9e8bd1ed8cdad82b5966c3e0ecae09fb1907d58efaa013543185d0a"
+
 [[package]]
 name = "serde"
 version = "1.0.228"

tools/virtual_dataset_generator/Cargo.toml

@@ -6,6 +6,7 @@ edition = "2021"
 [dependencies]
 anyhow = "1.0"
 clap = { version = "4.5", features = ["derive"] }
+regex = "1.11"
 rand = "0.8"
 rayon = "1.10"
 serde = { version = "1.0", features = ["derive"] }

tools/virtual_dataset_generator/src/bin/case_combo_generator.rs

@@ -0,0 +1,595 @@
+use anyhow::{Context, Result};
+use clap::Parser;
+use rand::rngs::StdRng;
+use rand::seq::SliceRandom;
+use rand::Rng;
+use rand::SeedableRng;
+use regex::Regex;
+use serde::{Deserialize, Serialize};
+use serde_json::Value;
+use std::collections::{HashMap, HashSet};
+use std::fs::File;
+use std::io::{BufRead, BufReader, BufWriter, Write};
+use std::path::PathBuf;
+
+#[derive(Parser, Debug)]
+#[command(
+    about = "Generate combo-heavy AniFileBERT char samples from real DMHY rows",
+    version
+)]
+struct Args {
+    #[arg(long)]
+    input: PathBuf,
+
+    #[arg(long)]
+    output: PathBuf,
+
+    #[arg(long, action = clap::ArgAction::Append)]
+    failure_report: Vec<PathBuf>,
+
+    #[arg(long, default_value_t = 12_000)]
+    selected_samples: usize,
+
+    #[arg(long, default_value_t = 60_000)]
+    context_samples: usize,
+
+    #[arg(long, default_value_t = 4)]
+    combos_per_row: usize,
+
+    #[arg(long, default_value_t = 42)]
+    seed: u64,
+}
+
+#[derive(Clone, Debug, Serialize, Deserialize)]
+struct CharRow {
+    filename: String,
+    tokens: Vec<String>,
+    labels: Vec<String>,
+    #[serde(default)]
+    tokenizer_variant: Option<String>,
+    #[serde(default)]
+    source: Option<String>,
+}
+
+fn main() -> Result<()> {
+    let args = Args::parse();
+    let target_re = Regex::new(
+        r"(?ix)(?:\b(?:S\d{1,2}|Season\s*\d{1,2}|第\s*\d+\s*[季期部章]|[ⅡⅢⅣⅤ]|II|III|IV|V)\b|\b(?:NCOP|NCED|OP|ED|PV|CM|TVCM|OVA|OAD|SP)\d*\b|(?:劇場版|特别篇|特別篇|ステージ|イベント|Event|Movie|PV|CM)\b|\b(?:Blu[-_ ]?ray\s*&\s*DVD|BD[-_ ]?BOX|Disc\.?\s*\d+|Vol\.?\s*\d+)\b|\b(?:h\.?264|x\.?264|h\.?265|x\.?265|AVC[-_ ]?YUV|yuv\d+p?\d*|AAC\([^)]*\))\b)",
+    )?;
+
+    let targets = failure_filenames(&args.failure_report)?;
+    let input =
+        File::open(&args.input).with_context(|| format!("open {}", args.input.display()))?;
+    let reader = BufReader::new(input);
+    let mut rng = StdRng::seed_from_u64(args.seed);
+
+    let mut selected_rows: Vec<CharRow> = Vec::new();
+    let mut context_rows: Vec<CharRow> = Vec::new();
+    let mut selected_seen = 0usize;
+    let mut context_seen = 0usize;
+    let mut seen_filenames = HashSet::new();
+    let mut total_rows = 0usize;
+
+    for line in reader.lines() {
+        let line = line?;
+        let line = line.trim();
+        if line.is_empty() {
+            continue;
+        }
+        total_rows += 1;
+        let row: CharRow = serde_json::from_str(line)
+            .with_context(|| format!("parse JSONL line {}", total_rows))?;
+        if row.filename.is_empty() || !seen_filenames.insert(row.filename.clone()) {
+            continue;
+        }
+        if targets.contains(&row.filename) || target_re.is_match(&row.filename) {
+            selected_seen += 1;
+            reservoir_push(
+                &mut selected_rows,
+                row,
+                args.selected_samples,
+                selected_seen,
+                &mut rng,
+            );
+        } else {
+            context_seen += 1;
+            reservoir_push(
+                &mut context_rows,
+                row,
+                args.context_samples,
+                context_seen,
+                &mut rng,
+            );
+        }
+    }
+
+    let mut out_rows: Vec<CharRow> = Vec::new();
+    let mut source_counts: HashMap<String, usize> = HashMap::new();
+
+    for row in &selected_rows {
+        let mut base = row.clone();
+        base.source = Some("combo_selected_base".to_string());
+        push_row(&mut out_rows, base, &mut source_counts);
+
+        let mut variants = build_combo_variants(row, 128);
+        variants.shuffle(&mut rng);
+        for variant in variants.into_iter().take(args.combos_per_row) {
+            push_row(&mut out_rows, variant, &mut source_counts);
+        }
+    }
+
+    for row in &context_rows {
+        let mut ctx = row.clone();
+        ctx.source = Some("combo_context".to_string());
+        push_row(&mut out_rows, ctx, &mut source_counts);
+    }
+
+    out_rows.shuffle(&mut rng);
+    let output =
+        File::create(&args.output).with_context(|| format!("create {}", args.output.display()))?;
+    let mut writer = BufWriter::new(output);
+    for row in &out_rows {
+        serde_json::to_writer(&mut writer, row)?;
+        writer.write_all(b"\n")?;
+    }
+    writer.flush()?;
+
+    let summary = serde_json::json!({
+        "input": args.input,
+        "output": args.output,
+        "total_rows": total_rows,
+        "failure_targets": targets.len(),
+        "selected_rows": selected_rows.len(),
+        "context_rows": context_rows.len(),
+        "written_rows": out_rows.len(),
+        "source_counts": source_counts,
+    });
+    println!("{}", serde_json::to_string_pretty(&summary)?);
+    Ok(())
+}
+
+fn reservoir_push<T: Clone>(
+    rows: &mut Vec<T>,
+    item: T,
+    limit: usize,
+    seen_count: usize,
+    rng: &mut StdRng,
+) {
+    if limit == 0 {
+        return;
+    }
+    if rows.len() < limit {
+        rows.push(item);
+        return;
+    }
+    let index = rng.gen_range(0..seen_count);
+    if index < limit {
+        rows[index] = item;
+    }
+}
+
+fn push_row(row: &mut Vec<CharRow>, item: CharRow, counts: &mut HashMap<String, usize>) {
+    let source = item.source.clone().unwrap_or_else(|| "unknown".to_string());
+    *counts.entry(source).or_insert(0) += 1;
+    row.push(item);
+}
+
+fn failure_filenames(report_paths: &[PathBuf]) -> Result<HashSet<String>> {
+    let mut filenames = HashSet::new();
+    for path in report_paths {
+        if !path.exists() {
+            continue;
+        }
+        let text = std::fs::read_to_string(path)
+            .with_context(|| format!("read failure report {}", path.display()))?;
+        let report: Value = serde_json::from_str(&text)
+            .with_context(|| format!("parse failure report {}", path.display()))?;
+        if let Some(modes) = report.get("modes").and_then(|v| v.as_object()) {
+            for mode in modes.values() {
+                if let Some(mode_obj) = mode.as_object() {
+                    if let Some(failures) = mode_obj.get("failures").and_then(|v| v.as_array()) {
+                        for failure in failures {
+                            if let Some(filename) = failure.get("filename").and_then(|v| v.as_str())
+                            {
+                                filenames.insert(filename.to_string());
+                            }
+                        }
+                    }
+                    if let Some(results) = mode_obj.get("results").and_then(|v| v.as_array()) {
+                        for result in results {
+                            if result.get("ok").and_then(|v| v.as_bool()) == Some(true) {
+                                continue;
+                            }
+                            if let Some(filename) = result.get("filename").and_then(|v| v.as_str())
+                            {
+                                filenames.insert(filename.to_string());
+                            }
+                        }
+                    }
+                }
+            }
+        }
+    }
+    Ok(filenames)
+}
+
+fn build_combo_variants(row: &CharRow, max_chars: usize) -> Vec<CharRow> {
+    let entities = extract_entities_from_labels(&row.tokens, &row.labels);
+    let title = first_value(&entities, "TITLE");
+    let season = first_value(&entities, "SEASON");
+    let episode = first_value(&entities, "EPISODE");
+    let special = first_value(&entities, "SPECIAL");
+    let resolution = first_value(&entities, "RESOLUTION");
+    let source = first_value(&entities, "SOURCE");
+
+    let mut specs: Vec<(String, Vec<(String, String)>, &'static str)> = Vec::new();
+    if let Some(title) = title.clone() {
+        specs.push((
+            title.clone(),
+            vec![(title.clone(), "TITLE".to_string())],
+            "combo_title",
+        ));
+        if let Some(season) = season.clone() {
+            specs.push((
+                format!("{title} {season}"),
+                vec![
+                    (title.clone(), "TITLE".to_string()),
+                    (season.clone(), "SEASON".to_string()),
+                ],
+                "combo_title_season",
+            ));
+            if let Some(episode) = episode.clone() {
+                specs.push((
+                    format!("{title} {season} {episode}"),
+                    vec![
+                        (title.clone(), "TITLE".to_string()),
+                        (season.clone(), "SEASON".to_string()),
+                        (episode.clone(), "EPISODE".to_string()),
+                    ],
+                    "combo_title_season_episode",
+                ));
+                if let (Some(resolution), Some(source)) = (resolution.clone(), source.clone()) {
+                    specs.push((
+                        format!("{title} {season} {episode} [{resolution}][{source}]"),
+                        vec![
+                            (title.clone(), "TITLE".to_string()),
+                            (season.clone(), "SEASON".to_string()),
+                            (episode.clone(), "EPISODE".to_string()),
+                            (resolution.clone(), "RESOLUTION".to_string()),
+                            (source.clone(), "SOURCE".to_string()),
+                        ],
+                        "combo_title_season_episode_resolution_source",
+                    ));
+                }
+            }
+        }
+    }
+    if let Some(episode) = episode.clone() {
+        if let Some(resolution) = resolution.clone() {
+            specs.push((
+                format!("{episode} [{resolution}]"),
+                vec![
+                    (episode.clone(), "EPISODE".to_string()),
+                    (resolution.clone(), "RESOLUTION".to_string()),
+                ],
+                "combo_episode_resolution",
+            ));
+            if let Some(source) = source.clone() {
+                specs.push((
+                    format!("{episode} [{resolution}][{source}]"),
+                    vec![
+                        (episode.clone(), "EPISODE".to_string()),
+                        (resolution.clone(), "RESOLUTION".to_string()),
+                        (source, "SOURCE".to_string()),
+                    ],
+                    "combo_episode_resolution_source",
+                ));
+            }
+        }
+    }
+    if let Some(special) = special.clone() {
+        specs.push((
+            special.clone(),
+            vec![(special.clone(), "SPECIAL".to_string())],
+            "combo_special_only",
+        ));
+    }
+    if let (Some(title), Some(special)) = (title.clone(), special.clone()) {
+        specs.push((
+            format!("{title} - {special}"),
+            vec![
+                (title.clone(), "TITLE".to_string()),
+                (special.clone(), "SPECIAL".to_string()),
+            ],
+            "combo_title_special",
+        ));
+        if let Some(episode) = episode.clone() {
+            specs.push((
+                format!("{title} - {special} [{episode}]"),
+                vec![
+                    (title.clone(), "TITLE".to_string()),
+                    (special.clone(), "SPECIAL".to_string()),
+                    (episode, "EPISODE".to_string()),
+                ],
+                "combo_title_special_episode",
+            ));
+        }
+        if let (Some(resolution), Some(source)) = (resolution.clone(), source.clone()) {
+            specs.push((
+                format!("{title} - {special} [{resolution}][{source}]"),
+                vec![
+                    (title, "TITLE".to_string()),
+                    (special, "SPECIAL".to_string()),
+                    (resolution.clone(), "RESOLUTION".to_string()),
+                    (source, "SOURCE".to_string()),
+                ],
+                "combo_title_special_resolution_source",
+            ));
+        }
+    }
+    if let (Some(title), Some(resolution), Some(source)) =
+        (title, resolution.clone(), source.clone())
+    {
+        specs.push((
+            format!("{title} [{resolution}][{source}]"),
+            vec![
+                (title.clone(), "TITLE".to_string()),
+                (resolution.clone(), "RESOLUTION".to_string()),
+                (source, "SOURCE".to_string()),
+            ],
+            "combo_title_resolution_source",
+        ));
+    }
+
+    let mut variants = Vec::new();
+    let mut seen_texts = HashSet::new();
+    for (text, spans, source_name) in specs {
+        if text.len() < 2 || text.len() > max_chars || !seen_texts.insert(text.clone()) {
+            continue;
+        }
+        if let Some(item) = char_item(&text, &spans, source_name) {
+            variants.push(item);
+        }
+    }
+    variants
+}
+
+fn extract_entities_from_labels(
+    tokens: &[String],
+    labels: &[String],
+) -> HashMap<String, Vec<String>> {
+    let mut entities: HashMap<String, Vec<String>> = HashMap::new();
+    let mut active_entity: Option<String> = None;
+    let mut active_tokens: Vec<String> = Vec::new();
+
+    for (token, label) in tokens.iter().zip(labels.iter()) {
+        if let Some(rest) = label.strip_prefix("B-") {
+            if let Some(entity) = active_entity.take() {
+                if !active_tokens.is_empty() {
+                    entities
+                        .entry(entity)
+                        .or_default()
+                        .push(active_tokens.join(""));
+                }
+            }
+            active_entity = Some(rest.to_string());
+            active_tokens = vec![token.clone()];
+        } else if let Some(rest) = label.strip_prefix("I-") {
+            if active_entity.as_deref() == Some(rest) {
+                active_tokens.push(token.clone());
+            } else {
+                if let Some(entity) = active_entity.take() {
+                    if !active_tokens.is_empty() {
+                        entities
+                            .entry(entity)
+                            .or_default()
+                            .push(active_tokens.join(""));
+                    }
+                }
+                active_entity = Some(rest.to_string());
+                active_tokens = vec![token.clone()];
+            }
+        } else {
+            if let Some(entity) = active_entity.take() {
+                if !active_tokens.is_empty() {
+                    entities
+                        .entry(entity)
+                        .or_default()
+                        .push(active_tokens.join(""));
+                }
+            }
+            active_tokens.clear();
+        }
+    }
+
+    if let Some(entity) = active_entity.take() {
+        if !active_tokens.is_empty() {
+            entities
+                .entry(entity)
+                .or_default()
+                .push(active_tokens.join(""));
+        }
+    }
+    entities
+}
+
+fn first_value(entities: &HashMap<String, Vec<String>>, name: &str) -> Option<String> {
+    entities
+        .get(name)
+        .and_then(|values| values.iter().find(|value| !value.trim().is_empty()))
+        .map(|value| value.trim().to_string())
+}
+
+fn char_item(filename: &str, spans: &[(String, String)], source: &str) -> Option<CharRow> {
+    let mut labels = vec!["O".to_string(); filename.chars().count()];
+    let tokens: Vec<String> = filename.chars().map(|c| c.to_string()).collect();
+    let mut cursor = 0usize;
+    for (text, entity) in spans {
+        if text.is_empty() {
+            continue;
+        }
+        if let Some(start) = find_substring(filename, text, cursor) {
+            let end = start + text.chars().count();
+            if start < labels.len() {
+                labels[start] = format!("B-{entity}");
+                for idx in (start + 1)..end.min(labels.len()) {
+                    labels[idx] = format!("I-{entity}");
+                }
+                cursor = end;
+                continue;
+            }
+        }
+        if let Some(start) = find_substring(filename, text, 0) {
+            let end = start + text.chars().count();
+            if start < labels.len() {
+                labels[start] = format!("B-{entity}");
+                for idx in (start + 1)..end.min(labels.len()) {
+                    labels[idx] = format!("I-{entity}");
+                }
+                cursor = end;
+                continue;
+            }
+        }
+        return None;
+    }
+
+    Some(CharRow {
+        filename: filename.to_string(),
+        tokens,
+        labels,
+        tokenizer_variant: Some("char".to_string()),
+        source: Some(source.to_string()),
+    })
+}
+
+fn find_substring(haystack: &str, needle: &str, start_char: usize) -> Option<usize> {
+    let mut char_indices: Vec<(usize, char)> = haystack.char_indices().collect();
+    char_indices.push((haystack.len(), '\0'));
+    let start_byte = if start_char >= char_indices.len() - 1 {
+        haystack.len()
+    } else {
+        char_indices[start_char].0
+    };
+    haystack[start_byte..].find(needle).map(|byte_offset| {
+        haystack[start_byte..start_byte + byte_offset]
+            .chars()
+            .count()
+            + start_char
+    })
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::fs;
+    use std::time::{SystemTime, UNIX_EPOCH};
+
+    fn unique_temp_path(prefix: &str) -> PathBuf {
+        let stamp = SystemTime::now()
+            .duration_since(UNIX_EPOCH)
+            .expect("clock drift")
+            .as_nanos();
+        std::env::temp_dir().join(format!("{prefix}_{stamp}_{}.json", std::process::id()))
+    }
+
+    fn make_row(filename: &str, spans: &[(String, String)], source: &str) -> CharRow {
+        char_item(filename, spans, source).expect("expected valid char row")
+    }
+
+    #[test]
+    fn failure_filenames_collects_failed_results_and_deduplicates() {
+        let path = unique_temp_path("case_combo_report");
+        let report = serde_json::json!({
+            "modes": {
+                "model_only": {
+                    "failures": [
+                        {"filename": "a.mkv"},
+                        {"filename": "a.mkv"}
+                    ],
+                    "results": [
+                        {"filename": "b.mkv", "ok": false},
+                        {"filename": "c.mkv", "ok": true}
+                    ]
+                },
+                "normalized_only": {
+                    "results": [
+                        {"filename": "d.mkv", "ok": false}
+                    ]
+                }
+            }
+        });
+        fs::write(&path, serde_json::to_string(&report).unwrap()).unwrap();
+
+        let set = failure_filenames(&[path.clone()]).unwrap();
+        assert!(set.contains("a.mkv"));
+        assert!(set.contains("b.mkv"));
+        assert!(set.contains("d.mkv"));
+        assert!(!set.contains("c.mkv"));
+        assert_eq!(set.len(), 3);
+
+        let _ = fs::remove_file(path);
+    }
+
+    #[test]
+    fn build_combo_variants_includes_full_episode_source_combo() {
+        let row = make_row(
+            "One Piece Season 21 1110 [1080p][WEB-DL].mkv",
+            &[
+                ("One Piece".to_string(), "TITLE".to_string()),
+                ("Season 21".to_string(), "SEASON".to_string()),
+                ("1110".to_string(), "EPISODE".to_string()),
+                ("1080p".to_string(), "RESOLUTION".to_string()),
+                ("WEB-DL".to_string(), "SOURCE".to_string()),
+            ],
+            "combo_seed",
+        );
+
+        let variants = build_combo_variants(&row, 128);
+        let combo = variants
+            .iter()
+            .find(|item| {
+                item.source.as_deref() == Some("combo_title_season_episode_resolution_source")
+            })
+            .expect("missing title+season+episode+resolution+source combo");
+        assert_eq!(combo.filename, "One Piece Season 21 1110 [1080p][WEB-DL]");
+        assert_eq!(
+            &combo.labels[0..9],
+            &[
+                "B-TITLE", "I-TITLE", "I-TITLE", "I-TITLE", "I-TITLE", "I-TITLE", "I-TITLE",
+                "I-TITLE", "I-TITLE"
+            ]
+        );
+        assert_eq!(
+            &combo.labels[10..19],
+            &[
+                "B-SEASON", "I-SEASON", "I-SEASON", "I-SEASON", "I-SEASON", "I-SEASON", "I-SEASON",
+                "I-SEASON", "I-SEASON"
+            ]
+        );
+        assert_eq!(
+            &combo.labels[20..24],
+            &["B-EPISODE", "I-EPISODE", "I-EPISODE", "I-EPISODE"]
+        );
+        assert_eq!(
+            &combo.labels[26..31],
+            &[
+                "B-RESOLUTION",
+                "I-RESOLUTION",
+                "I-RESOLUTION",
+                "I-RESOLUTION",
+                "I-RESOLUTION"
+            ]
+        );
+        assert_eq!(
+            &combo.labels[33..39],
+            &["B-SOURCE", "I-SOURCE", "I-SOURCE", "I-SOURCE", "I-SOURCE", "I-SOURCE"]
+        );
+        assert_eq!(combo.labels[9], "O");
+        assert_eq!(combo.labels[19], "O");
+        assert_eq!(combo.labels[24], "O");
+        assert_eq!(combo.labels[25], "O");
+        assert_eq!(combo.labels[31], "O");
+        assert_eq!(combo.labels[32], "O");
+        assert_eq!(combo.labels[39], "O");
+    }
+}