use anyhow::{bail, Context, Result}; use clap::Parser; use rayon::ThreadPoolBuilder; use regex::Regex; use serde::{Deserialize, Serialize}; #[cfg(test)] use serde_json::Value; use std::collections::{HashMap, HashSet}; use std::fs::{self, File}; use std::io::{BufRead, BufReader, BufWriter, Write}; use std::path::{Path, PathBuf}; use std::sync::OnceLock; const DEFAULT_MAX_ROWS: usize = 50_000; const NUMERIC_RATIO_PERCENT: usize = 60; const DEFAULT_PATH_SPECIAL_RATIO_PERCENT: usize = 25; #[derive(Parser, Debug)] #[command(about = "Generate schema v2 numeric-title and path synthetic augmentation JSONL")] struct Args { #[arg( long, default_value = "reports/dmhy_template_recipes.full_top5000.seed.jsonl" )] recipes: PathBuf, #[arg(long, default_value = "label_schema.json")] label_schema_file: PathBuf, #[arg(long, default_value = "data/synthetic_numeric_titles.txt")] numeric_title_seeds: PathBuf, #[arg(long, default_value = "data/synthetic_path_prefixes.txt")] path_prefix_seeds: PathBuf, #[arg(long, default_value_t = 3000)] limit_templates: usize, #[arg(long, default_value_t = DEFAULT_MAX_ROWS)] max_rows: usize, #[arg(long, default_value_t = DEFAULT_PATH_SPECIAL_RATIO_PERCENT)] path_special_ratio_percent: usize, #[arg(long, default_value = "data/schema_v2_synthetic_aug.jsonl")] output: PathBuf, #[arg(long, default_value = "data/schema_v2_synthetic_aug.manifest.json")] manifest_output: PathBuf, #[arg(long, default_value_t = 0)] threads: usize, } #[derive(Debug, Deserialize)] struct LabelSchema { labels: Vec, } #[derive(Debug, Clone, Deserialize)] struct TitleSpan { role_indices: Vec, } #[derive(Debug, Clone, Deserialize)] struct Recipe { template_id: String, template: String, roles: Vec, #[serde(default)] title_spans: Vec, confidence: Option, #[serde(rename = "count")] _count: Option, } #[derive(Debug, Clone, Serialize)] struct Record { filename: String, tokens: Vec, labels: Vec, source: String, template_id: String, template: String, synthetic_kind: String, tokenizer_variant: String, } #[derive(Debug, Default, Serialize)] struct Manifest { generated_rows: usize, dropped_conflicting_templates: usize, deduped_rows: usize, numeric_title_rows: usize, path_rows: usize, path_series_rows: usize, path_movie_rows: usize, path_special_rows: usize, path_confuser_rows: usize, dropped_media_kind_mismatch: usize, recipes: String, label_schema_file: String, numeric_title_seeds: String, path_prefix_seeds: String, limit_templates: usize, max_rows: usize, templates_loaded: usize, numeric_seed_count: usize, path_seed_count: usize, path_series_seed_count: usize, path_movie_seed_count: usize, path_special_seed_count: usize, path_confuser_seed_count: usize, } #[derive(Debug, Default)] struct BuildStats { dropped_conflicting_templates: usize, deduped_rows: usize, numeric_title_rows: usize, path_rows: usize, path_series_rows: usize, path_movie_rows: usize, path_special_rows: usize, path_confuser_rows: usize, dropped_media_kind_mismatch: usize, } #[derive(Debug, Clone, Copy, Eq, PartialEq)] enum PathSeedKind { Series, Movie, Special, Confuser, } impl PathSeedKind { const ALL: [PathSeedKind; 4] = [ PathSeedKind::Series, PathSeedKind::Movie, PathSeedKind::Special, PathSeedKind::Confuser, ]; fn synthetic_kind(self) -> &'static str { match self { PathSeedKind::Series => "path_series", PathSeedKind::Movie => "path_movie", PathSeedKind::Special => "path_special", PathSeedKind::Confuser => "path_confuser", } } fn source(self) -> &'static str { match self { PathSeedKind::Series => "schema_v2_synthetic_path_series", PathSeedKind::Movie => "schema_v2_synthetic_path_movie", PathSeedKind::Special => "schema_v2_synthetic_path_special", PathSeedKind::Confuser => "schema_v2_synthetic_path_confuser", } } } #[derive(Debug, Clone, Copy, Eq, PartialEq)] enum BaseMediaKind { SeriesEpisode, Movie, Special, Unknown, } #[derive(Debug, Clone)] struct PathSeed { template: String, kind: PathSeedKind, } #[derive(Debug, Clone, Copy, Eq, PartialEq)] enum Entity { TitleLatin, PathTitleLatin, PathSeason, Season, Episode, Special, Group, Resolution, Source, Tag, } impl Entity { fn b_label(self) -> &'static str { match self { Entity::TitleLatin => "B-TITLE_LATIN", Entity::PathTitleLatin => "B-PATH_TITLE_LATIN", Entity::PathSeason => "B-PATH_SEASON", Entity::Season => "B-SEASON", Entity::Episode => "B-EPISODE", Entity::Special => "B-SPECIAL", Entity::Group => "B-GROUP", Entity::Resolution => "B-RESOLUTION", Entity::Source => "B-SOURCE", Entity::Tag => "B-TAG", } } fn i_label(self) -> &'static str { match self { Entity::TitleLatin => "I-TITLE_LATIN", Entity::PathTitleLatin => "I-PATH_TITLE_LATIN", Entity::PathSeason => "I-PATH_SEASON", Entity::Season => "I-SEASON", Entity::Episode => "I-EPISODE", Entity::Special => "I-SPECIAL", Entity::Group => "I-GROUP", Entity::Resolution => "I-RESOLUTION", Entity::Source => "I-SOURCE", Entity::Tag => "I-TAG", } } } #[derive(Debug, Default, Clone)] struct CharBuilder { filename: String, tokens: Vec, labels: Vec, } impl CharBuilder { fn append_o(&mut self, text: &str) { self.append(text, None); } fn append_entity(&mut self, text: &str, entity: Entity) { self.append(text, Some(entity)); } fn append(&mut self, text: &str, entity: Option) { let mut first = true; for ch in text.chars() { self.filename.push(ch); self.tokens.push(ch.to_string()); let label = match entity { None => "O", Some(value) => { if first { value.b_label() } else { value.i_label() } } }; self.labels.push(label.to_string()); first = false; } } fn into_parts(self) -> (String, Vec, Vec) { (self.filename, self.tokens, self.labels) } } fn main() -> Result<()> { let args = Args::parse(); if args.max_rows == 0 { bail!("--max-rows must be positive"); } if args.threads > 0 { ThreadPoolBuilder::new() .num_threads(args.threads) .build_global() .context("failed to configure rayon thread pool")?; } let label_set = load_label_set(&args.label_schema_file)?; let recipes = load_recipes(&args.recipes, args.limit_templates)?; let numeric_titles = load_seed_lines(&args.numeric_title_seeds)?; let path_prefixes = load_path_seed_lines(&args.path_prefix_seeds)?; if numeric_titles.is_empty() { bail!("numeric title seed file is empty"); } if path_prefixes.is_empty() { bail!("path prefix seed file is empty"); } let mut records = Vec::with_capacity(args.max_rows); let mut seen = HashSet::new(); let mut stats = BuildStats::default(); let numeric_quota = (args.max_rows * NUMERIC_RATIO_PERCENT / 100).max(1); let path_quota = args.max_rows.saturating_sub(numeric_quota).max(1); let special_ratio = args.path_special_ratio_percent.min(100); let mut path_special_quota = path_quota * special_ratio / 100; let mut path_series_quota = path_quota.saturating_sub(path_special_quota); let path_series_seed_count = path_seed_count(&path_prefixes, PathSeedKind::Series); let path_movie_seed_count = path_seed_count(&path_prefixes, PathSeedKind::Movie); let path_special_seed_count = path_seed_count(&path_prefixes, PathSeedKind::Special); let path_confuser_seed_count = path_seed_count(&path_prefixes, PathSeedKind::Confuser); let non_series_kind_count = [ path_movie_seed_count, path_special_seed_count, path_confuser_seed_count, ] .into_iter() .filter(|count| *count > 0) .count(); if path_series_seed_count == 0 { path_series_quota = 0; path_special_quota = path_quota; } else if non_series_kind_count == 0 { path_series_quota = path_quota; path_special_quota = 0; } for recipe in &recipes { if stats.numeric_title_rows >= numeric_quota { break; } if !recipe.roles.iter().any(|role| role == "TITLE") { stats.dropped_conflicting_templates += 1; continue; } for (seed_idx, title) in numeric_titles.iter().enumerate() { if stats.numeric_title_rows >= numeric_quota { break; } match build_numeric_record(recipe, title, seed_idx) { Some(record) => { push_record(record, &mut records, &mut seen, &mut stats, &label_set)? } None => stats.dropped_conflicting_templates += 1, } } } let base_files = path_base_records(&records); 'path_kind: for kind in PathSeedKind::ALL { let seed_count = path_seed_count(&path_prefixes, kind); if seed_count == 0 { continue; } let kind_quota = path_kind_quota( kind, path_series_quota, path_special_quota, non_series_kind_count, ); if kind_quota == 0 { continue; } for (base_idx, base) in base_files.iter().enumerate() { for (prefix_idx, prefix) in path_prefixes .iter() .filter(|seed| seed.kind == kind) .enumerate() { if path_kind_rows(&stats, kind) >= kind_quota || stats.path_rows >= path_quota { continue 'path_kind; } let title = path_title_for(prefix_idx + base_idx, &numeric_titles); let season_dir = path_season_dir_for(prefix_idx + base_idx); match build_path_record(&prefix.template, kind, &title, season_dir, base) { Some(record) => { push_record(record, &mut records, &mut seen, &mut stats, &label_set)?; } None => { if !path_seed_compatible_with_base(kind, classify_base_media_kind(base)) { stats.dropped_media_kind_mismatch += 1; } else { stats.dropped_conflicting_templates += 1; } } } } } } if records.is_empty() { bail!("no synthetic rows were generated"); } if let Some(parent) = args.output.parent() { fs::create_dir_all(parent) .with_context(|| format!("failed to create output directory {}", parent.display()))?; } let output = File::create(&args.output) .with_context(|| format!("failed to create {}", args.output.display()))?; let mut writer = BufWriter::new(output); for record in &records { serde_json::to_writer(&mut writer, record)?; writer.write_all(b"\n")?; } if let Some(parent) = args.manifest_output.parent() { fs::create_dir_all(parent) .with_context(|| format!("failed to create manifest directory {}", parent.display()))?; } let manifest = Manifest { generated_rows: records.len(), dropped_conflicting_templates: stats.dropped_conflicting_templates, deduped_rows: stats.deduped_rows, numeric_title_rows: stats.numeric_title_rows, path_rows: stats.path_rows, path_series_rows: stats.path_series_rows, path_movie_rows: stats.path_movie_rows, path_special_rows: stats.path_special_rows, path_confuser_rows: stats.path_confuser_rows, dropped_media_kind_mismatch: stats.dropped_media_kind_mismatch, recipes: args.recipes.display().to_string(), label_schema_file: args.label_schema_file.display().to_string(), numeric_title_seeds: args.numeric_title_seeds.display().to_string(), path_prefix_seeds: args.path_prefix_seeds.display().to_string(), limit_templates: args.limit_templates, max_rows: args.max_rows, templates_loaded: recipes.len(), numeric_seed_count: numeric_titles.len(), path_seed_count: path_prefixes.len(), path_series_seed_count, path_movie_seed_count, path_special_seed_count, path_confuser_seed_count, }; fs::write( &args.manifest_output, serde_json::to_string_pretty(&manifest)?, ) .with_context(|| format!("failed to write {}", args.manifest_output.display()))?; println!("{}", serde_json::to_string_pretty(&manifest)?); Ok(()) } fn load_label_set(path: &Path) -> Result> { let text = fs::read_to_string(path).with_context(|| format!("failed to read {}", path.display()))?; let schema: LabelSchema = serde_json::from_str(&text).with_context(|| format!("invalid {}", path.display()))?; Ok(schema.labels.into_iter().collect()) } fn load_seed_lines(path: &Path) -> Result> { let file = File::open(path).with_context(|| format!("failed to open {}", path.display()))?; let reader = BufReader::new(file); let mut lines = Vec::new(); let mut seen = HashSet::new(); for line in reader.lines() { let line = line?; let Some(seed) = parse_seed_line(&line) else { continue; }; if seen.insert(seed.clone()) { lines.push(seed); } } Ok(lines) } fn load_path_seed_lines(path: &Path) -> Result> { let file = File::open(path).with_context(|| format!("failed to open {}", path.display()))?; let reader = BufReader::new(file); let mut lines = Vec::new(); let mut seen = HashSet::new(); for line in reader.lines() { let line = line?; let Some(seed) = parse_path_seed_line(&line) else { continue; }; let key = format!("{}|{}", seed.kind.synthetic_kind(), seed.template); if seen.insert(key) { lines.push(seed); } } Ok(lines) } fn parse_path_seed_line(line: &str) -> Option { let mut trimmed = line.trim(); if trimmed.is_empty() { return None; } if let Some(comment_body) = trimmed.strip_prefix('#') { trimmed = comment_body.trim(); if !trimmed.contains('|') { return None; } } let mut kind = PathSeedKind::Series; if let Some((kind_text, value)) = trimmed.split_once('|') { match kind_text.trim() { "series" | "path" => { kind = PathSeedKind::Series; trimmed = value.trim(); } "movie" | "film" | "theatrical" => { kind = PathSeedKind::Movie; trimmed = value.trim(); } "special" => { kind = PathSeedKind::Special; trimmed = value.trim(); } "confuser" => { kind = PathSeedKind::Confuser; trimmed = value.trim(); } _ => {} } } if trimmed.is_empty() || trimmed.starts_with('#') { None } else { Some(PathSeed { template: trimmed.to_string(), kind, }) } } fn parse_seed_line(line: &str) -> Option { let mut trimmed = line.trim(); if trimmed.is_empty() { return None; } if let Some(comment_body) = trimmed.strip_prefix('#') { trimmed = comment_body.trim(); if !trimmed.contains('|') { return None; } } if let Some((kind, value)) = trimmed.split_once('|') { let known_kind = matches!( kind.trim(), "series" | "movie" | "film" | "theatrical" | "special" | "confuser" | "path" | "numeric" ); if known_kind { trimmed = value.trim(); } } if trimmed.is_empty() || trimmed.starts_with('#') { None } else { Some(trimmed.to_string()) } } fn path_seed_count(path_prefixes: &[PathSeed], kind: PathSeedKind) -> usize { path_prefixes .iter() .filter(|seed| seed.kind == kind) .count() } fn path_kind_quota( kind: PathSeedKind, series_quota: usize, non_series_quota: usize, non_series_kind_count: usize, ) -> usize { match kind { PathSeedKind::Series => series_quota, PathSeedKind::Movie | PathSeedKind::Special | PathSeedKind::Confuser => { if non_series_kind_count == 0 { 0 } else { (non_series_quota / non_series_kind_count).max(1) } } } } fn path_kind_rows(stats: &BuildStats, kind: PathSeedKind) -> usize { match kind { PathSeedKind::Series => stats.path_series_rows, PathSeedKind::Movie => stats.path_movie_rows, PathSeedKind::Special => stats.path_special_rows, PathSeedKind::Confuser => stats.path_confuser_rows, } } fn load_recipes(path: &Path, limit_templates: usize) -> Result> { let file = File::open(path).with_context(|| format!("failed to open {}", path.display()))?; let reader = BufReader::new(file); let mut recipes = Vec::new(); for (idx, line) in reader.lines().enumerate() { if recipes.len() >= limit_templates { break; } let line = line.with_context(|| format!("failed reading recipe line {}", idx + 1))?; if line.trim().is_empty() { continue; } let recipe: Recipe = serde_json::from_str(&line) .with_context(|| format!("invalid recipe JSONL line {}", idx + 1))?; if recipe.confidence.as_deref() != Some("high") { continue; } if recipe.template.split_whitespace().count() != recipe.roles.len() { continue; } recipes.push(recipe); } Ok(recipes) } fn push_record( record: Record, records: &mut Vec, seen: &mut HashSet, stats: &mut BuildStats, label_set: &HashSet, ) -> Result<()> { validate_record_labels(&record, label_set)?; let key = format!("{}\u{1f}{}", record.filename, record.labels.join("\u{1f}")); if !seen.insert(key) { stats.deduped_rows += 1; return Ok(()); } match record.synthetic_kind.as_str() { "numeric_title" => stats.numeric_title_rows += 1, "path_series" => { stats.path_rows += 1; stats.path_series_rows += 1; } "path_movie" => { stats.path_rows += 1; stats.path_movie_rows += 1; } "path_special" => { stats.path_rows += 1; stats.path_special_rows += 1; } "path_confuser" => { stats.path_rows += 1; stats.path_confuser_rows += 1; } other => bail!("unknown synthetic_kind {other}"), } records.push(record); Ok(()) } fn validate_record_labels(record: &Record, label_set: &HashSet) -> Result<()> { if record.tokens.len() != record.labels.len() { bail!("tokens/labels mismatch for {}", record.filename); } if record.tokens != record .filename .chars() .map(|ch| ch.to_string()) .collect::>() { bail!("tokens are not char-level for {}", record.filename); } let mut prev_entity: Option<&str> = None; let mut has_title = false; for (idx, label) in record.labels.iter().enumerate() { if !label_set.contains(label) { bail!("label {label} is not present in label schema"); } if label == "O" { prev_entity = None; continue; } let Some((prefix, entity)) = label.split_once('-') else { bail!( "malformed BIO label {label} at {idx} for {}", record.filename ); }; if prefix != "B" && prefix != "I" { bail!( "invalid BIO prefix {prefix} at {idx} for {}", record.filename ); } if entity.is_empty() { bail!("empty BIO entity at {idx} for {}", record.filename); } if prefix == "I" && prev_entity != Some(entity) { bail!("orphan I-tag {label} at {idx} for {}", record.filename); } if entity.contains("TITLE") { has_title = true; } prev_entity = Some(entity); } if !has_title { bail!("missing title span for {}", record.filename); } Ok(()) } fn logical_title_spans(recipe: &Recipe) -> Option>> { if recipe.title_spans.is_empty() { return Some( recipe .roles .iter() .enumerate() .filter_map(|(index, role)| { if role == "TITLE" { Some(vec![index]) } else { None } }) .collect(), ); } let mut spans = Vec::new(); for span in &recipe.title_spans { if span.role_indices.is_empty() { return None; } for &index in &span.role_indices { if recipe.roles.get(index).map(String::as_str) != Some("TITLE") { return None; } } spans.push(span.role_indices.clone()); } Some(spans) } fn split_title_for_slots(title: &str, slots: usize) -> Option> { if slots == 0 { return None; } if slots == 1 { return Some(vec![title.to_string()]); } let words: Vec<&str> = title .split_whitespace() .filter(|word| !word.is_empty()) .collect(); if words.len() < slots { return None; } let mut chunks = Vec::with_capacity(slots); let mut start = 0usize; for slot in 0..slots { let remaining_words = words.len() - start; let remaining_slots = slots - slot; let take = remaining_words.div_ceil(remaining_slots); let end = start + take; chunks.push(words[start..end].join(" ")); start = end; } Some(chunks) } fn build_numeric_record(recipe: &Recipe, title: &str, variant: usize) -> Option { let classes: Vec<&str> = recipe.template.split_whitespace().collect(); if classes.len() != recipe.roles.len() { return None; } let title_spans = logical_title_spans(recipe)?; if title_spans.len() != 1 { return None; } let title_role_indices = &title_spans[0]; let title_chunks = split_title_for_slots(title, title_role_indices.len())?; let title_chunk_by_index: HashMap = title_role_indices .iter() .copied() .zip(title_chunks.iter().map(String::as_str)) .collect(); let mut builder = CharBuilder::default(); let mut previous_role = ""; for (index, (class_name, role)) in classes.iter().zip(recipe.roles.iter()).enumerate() { let special_number = role == "EPISODE" && previous_role == "SPECIAL"; let role_for_label = if special_number { "SPECIAL" } else { role.as_str() }; append_template_group( &mut builder, class_name, role, role_for_label, title_chunk_by_index.get(&index).copied().unwrap_or(title), variant, special_number, )?; if role != "O" { previous_role = role; } } let (filename, tokens, labels) = builder.into_parts(); if has_entity(&labels, "EPISODE") && title_has_episode_label(&filename, &labels, title) { return None; } Some(Record { filename, tokens, labels, source: "schema_v2_synthetic_numeric_title".to_string(), template_id: recipe.template_id.clone(), template: recipe.template.clone(), synthetic_kind: "numeric_title".to_string(), tokenizer_variant: "char".to_string(), }) } fn append_template_group( builder: &mut CharBuilder, class_name: &str, content_role: &str, role: &str, title: &str, variant: usize, special_number: bool, ) -> Option<()> { if class_name == "SEP" { builder.append_o(separator_for(variant)); return Some(()); } if class_name == "SXE" { append_sxe(builder); return Some(()); } let entity = entity_for_role(role, class_name); let content = if special_number { "01".to_string() } else { content_for_role(content_role, class_name, title, variant) }; if class_name.starts_with("BRACKET_") { builder.append_o("["); builder.append(&content, entity); builder.append_o("]"); } else { builder.append(&content, entity); } Some(()) } fn separator_for(variant: usize) -> &'static str { match variant % 5 { 0 => " ", 1 => " - ", 2 => ".", 3 => "_", _ => " ", } } fn entity_for_role(role: &str, class_name: &str) -> Option { match role { "TITLE" => Some(Entity::TitleLatin), "GROUP" => Some(Entity::Group), "SEASON" => Some(Entity::Season), "EPISODE" | "EPISODE_VERSION" | "EPISODE_RANGE" => Some(Entity::Episode), "SPECIAL" | "VOLUME" => Some(Entity::Special), "RESOLUTION" => Some(Entity::Resolution), "SOURCE" | "HASH" => Some(Entity::Source), "TAG" => Some(Entity::Tag), "O" => { if class_name.contains("HASH") { Some(Entity::Source) } else { None } } _ => None, } } fn content_for_role(role: &str, class_name: &str, title: &str, variant: usize) -> String { match role { "TITLE" => title.to_string(), "GROUP" => group_for(variant).to_string(), "SEASON" => season_for(variant).to_string(), "EPISODE" => episode_for(variant).to_string(), "EPISODE_VERSION" => format!("{}v2", episode_for(variant)), "EPISODE_RANGE" => "01-12".to_string(), "SPECIAL" => special_for(variant).to_string(), "VOLUME" => format!("Vol.{}", (variant % 6) + 1), "RESOLUTION" => resolution_for(variant).to_string(), "SOURCE" => source_for(class_name, variant).to_string(), "HASH" => "A1B2C3D4".to_string(), "TAG" => tag_for(variant).to_string(), _ => neutral_for_class(class_name, variant).to_string(), } } fn group_for(variant: usize) -> &'static str { const GROUPS: [&str; 8] = [ "DBD-Raws", "VCB-Studio", "A.I.Raws", "Lilith-Raws", "ANi", "Nekomoe", "Baha", "Skytree", ]; GROUPS[variant % GROUPS.len()] } fn season_for(variant: usize) -> &'static str { const SEASONS: [&str; 6] = ["S2", "Season 2", "S01", "2nd Season", "第2季", "Part 2"]; SEASONS[variant % SEASONS.len()] } fn episode_for(variant: usize) -> &'static str { const EPISODES: [&str; 8] = ["01", "02", "03", "12", "24", "087", "100", "S01E03"]; EPISODES[variant % EPISODES.len()] } fn special_for(variant: usize) -> &'static str { const SPECIALS: [&str; 12] = [ "NCED", "NCOP2", "PV", "CM", "Menu", "Trailer", "SP", "SP01", "OVA", "OAD", "OP2", "ED1", ]; SPECIALS[variant % SPECIALS.len()] } fn resolution_for(variant: usize) -> &'static str { const RESOLUTIONS: [&str; 5] = ["1080p", "720P", "2160p", "1920x1080", "4K"]; RESOLUTIONS[variant % RESOLUTIONS.len()] } fn source_for(class_name: &str, variant: usize) -> &'static str { if class_name.contains("MEDIA_BLOCK") { return "WEB-DL 1080p AVC AAC"; } const SOURCES: [&str; 10] = [ "WEB-DL", "BDRip", "HEVC", "AAC", "FLAC", "CHS", "JPSC", "NF", "Baha", "x264", ]; SOURCES[variant % SOURCES.len()] } fn tag_for(variant: usize) -> &'static str { const TAGS: [&str; 6] = ["Gekijouban", "Movie", "TV", "2004", "Extras", "SPs"]; TAGS[variant % TAGS.len()] } fn neutral_for_class(class_name: &str, variant: usize) -> &'static str { if class_name.contains("RESOLUTION") { resolution_for(variant) } else if class_name.contains("MEDIA") || class_name.contains("LANG") { source_for(class_name, variant) } else if class_name.contains("SPECIAL") { special_for(variant) } else if class_name.contains("EPISODE") { episode_for(variant) } else if class_name.contains("SEASON") { season_for(variant) } else if class_name.contains("HASH") { "A1B2C3D4" } else if class_name.contains("DATE") { "2024" } else { "v2" } } fn append_sxe(builder: &mut CharBuilder) { builder.append_o("S"); builder.append_entity("01", Entity::Season); builder.append_o("E"); builder.append_entity("03", Entity::Episode); } fn title_has_episode_label(filename: &str, labels: &[String], title: &str) -> bool { let Some(start) = filename.find(title) else { return false; }; let end = start + title.len(); let char_starts = char_byte_offsets(filename); char_starts .iter() .enumerate() .filter(|(_, byte_idx)| **byte_idx >= start && **byte_idx < end) .any(|(idx, _)| labels[idx].contains("EPISODE") || labels[idx].contains("SEASON")) } fn char_byte_offsets(text: &str) -> Vec { text.char_indices().map(|(idx, _)| idx).collect() } fn has_entity(labels: &[String], entity: &str) -> bool { labels.iter().any(|label| { label .strip_prefix("B-") .is_some_and(|value| value == entity) }) } fn path_base_records(records: &[Record]) -> Vec { let mut bases = Vec::new(); bases.extend(simple_path_leaf_records()); bases.extend( records .iter() .filter(|record| record.synthetic_kind == "numeric_title") .take(256) .cloned(), ); bases } fn classify_base_media_kind(record: &Record) -> BaseMediaKind { let has_episode = has_entity(&record.labels, "EPISODE"); let has_special = has_entity(&record.labels, "SPECIAL"); if has_special { return BaseMediaKind::Special; } if has_episode { return BaseMediaKind::SeriesEpisode; } let lower = record.filename.to_ascii_lowercase(); if ["movie", "gekijouban", "the movie", "film", "劇場版"] .iter() .any(|needle| lower.contains(&needle.to_ascii_lowercase())) { return BaseMediaKind::Movie; } BaseMediaKind::Unknown } fn path_seed_compatible_with_base(kind: PathSeedKind, base_kind: BaseMediaKind) -> bool { match kind { PathSeedKind::Series => base_kind == BaseMediaKind::SeriesEpisode, PathSeedKind::Movie => base_kind == BaseMediaKind::Movie, PathSeedKind::Special => base_kind == BaseMediaKind::Special, PathSeedKind::Confuser => true, } } fn simple_path_leaf_records() -> Vec { let mut records = Vec::new(); for episode in 1..=120 { let episode_text = format!("{episode:02}"); let filename = format!("{episode_text}.mkv"); let template_id = format!("path_leaf_episode_{episode_text}"); records.push(char_record_from_spans( &filename, &[(0, episode_text.len(), Entity::Episode)], "schema_v2_synthetic_path_leaf", &template_id, "path_leaf_episode", "path_aug", )); } records.extend([ char_record_from_spans( "file.mkv", &[], "schema_v2_synthetic_path_leaf", "path_leaf_file", "path_leaf_file", "path_aug", ), char_record_from_spans( "S01E03.mkv", &[(1, 3, Entity::Season), (4, 6, Entity::Episode)], "schema_v2_synthetic_path_leaf", "path_leaf_sxe", "path_leaf_sxe", "path_aug", ), ]); const MOVIE_TAGS: [&str; 6] = [ "Movie", "Gekijouban", "The Movie", "Film", "BDMovie", "劇場版", ]; const MOVIE_SUFFIXES: [&str; 12] = [ "", " 01", " 02", " 2004", " 1080p", " BDRip", " WEB-DL", " Remux", " Complete", " Director Cut", " Part A", " Part B", ]; for (tag_idx, tag) in MOVIE_TAGS.iter().enumerate() { for (suffix_idx, suffix) in MOVIE_SUFFIXES.iter().enumerate() { let filename = format!("{tag}{suffix}.mkv"); let template_id = format!("path_leaf_movie_{tag_idx}_{suffix_idx}"); records.push(char_record_from_spans( &filename, &[(0, tag.len(), Entity::Tag)], "schema_v2_synthetic_path_leaf", &template_id, "path_leaf_movie", "path_aug", )); } } const SPECIAL_TAGS: [&str; 12] = [ "NCOP", "NCOP2", "NCED", "NCED2", "OP", "OP2", "ED", "ED1", "PV", "CM", "Menu", "Trailer", ]; const SPECIAL_SUFFIXES: [&str; 8] = ["", " 01", " 02", " A", " B", " Clean", " Creditless", " v2"]; for (tag_idx, tag) in SPECIAL_TAGS.iter().enumerate() { for (suffix_idx, suffix) in SPECIAL_SUFFIXES.iter().enumerate() { let special_text = format!("{tag}{suffix}"); let filename = format!("{special_text}.mkv"); let template_id = format!("path_leaf_special_{tag_idx}_{suffix_idx}"); records.push(char_record_from_spans( &filename, &[(0, special_text.len(), Entity::Special)], "schema_v2_synthetic_path_leaf", &template_id, "path_leaf_special", "path_aug", )); } } records } fn char_record_from_spans( filename: &str, spans: &[(usize, usize, Entity)], source: &str, template_id: &str, template: &str, synthetic_kind: &str, ) -> Record { let mut char_entities: Vec> = vec![None; filename.chars().count()]; let byte_offsets = char_byte_offsets(filename); for (start, end, entity) in spans { for (char_idx, byte_idx) in byte_offsets.iter().enumerate() { if byte_idx >= start && byte_idx < end { char_entities[char_idx] = Some(*entity); } } } let mut tokens = Vec::new(); let mut labels = Vec::new(); let mut active_entity = None; for (ch, entity) in filename.chars().zip(char_entities) { tokens.push(ch.to_string()); let label = match entity { None => { active_entity = None; "O".to_string() } Some(value) => { let prefix = if active_entity == Some(value) { "I" } else { "B" }; active_entity = Some(value); format!( "{}-{}", prefix, value .b_label() .strip_prefix("B-") .expect("entity label has B- prefix") ) } }; labels.push(label); } Record { filename: filename.to_string(), tokens, labels, source: source.to_string(), template_id: template_id.to_string(), template: template.to_string(), synthetic_kind: synthetic_kind.to_string(), tokenizer_variant: "char".to_string(), } } fn path_title_for(index: usize, numeric_titles: &[String]) -> String { const GENERIC: [&str; 8] = [ "Title", "Naruto", "Sousou no Frieren", "Generic 2001 Story", "Yamada-kun to 7-nin no Majo", "91 Days", "Area 88", "No.6", ]; if index % 3 == 0 { GENERIC[index % GENERIC.len()].to_string() } else { numeric_titles[index % numeric_titles.len()].clone() } } fn path_season_dir_for(index: usize) -> &'static str { const SEASON_DIRS: [&str; 10] = [ "Season 01", "S01", "01", "第2季", "Season 2", "S2", "Gekijouban", "Movie", "2004", "TV", ]; SEASON_DIRS[index % SEASON_DIRS.len()] } fn build_path_record( prefix: &str, kind: PathSeedKind, title: &str, season_dir: &str, base: &Record, ) -> Option { if !prefix.contains("{filename}") { return None; } if !path_seed_compatible_with_base(kind, classify_base_media_kind(base)) { return None; } let mut builder = CharBuilder::default(); let mut remaining = prefix; while !remaining.is_empty() { if let Some(rest) = remaining.strip_prefix("{title}") { builder.append_entity(title, Entity::PathTitleLatin); remaining = rest; } else if let Some(rest) = remaining.strip_prefix("{filename}") { append_existing_record(&mut builder, base); remaining = rest; } else if let Some(rest) = remaining.strip_prefix("{season_dir}") { append_path_segment(&mut builder, season_dir); remaining = rest; } else { let next = next_placeholder_index(remaining).unwrap_or(remaining.len()); let literal = &remaining[..next]; append_path_literal(&mut builder, literal); remaining = &remaining[next..]; } } let (filename, tokens, labels) = builder.into_parts(); Some(Record { filename, tokens, labels, source: kind.source().to_string(), template_id: format!("path::{}", base.template_id), template: prefix.to_string(), synthetic_kind: kind.synthetic_kind().to_string(), tokenizer_variant: "char".to_string(), }) } fn next_placeholder_index(text: &str) -> Option { ["{title}", "{filename}", "{season_dir}"] .iter() .filter_map(|needle| text.find(needle)) .min() } fn append_existing_record(builder: &mut CharBuilder, record: &Record) { for (token, label) in record.tokens.iter().zip(record.labels.iter()) { builder.filename.push_str(token); builder.tokens.push(token.clone()); builder.labels.push(label.clone()); } } fn append_path_literal(builder: &mut CharBuilder, literal: &str) { let mut segment = String::new(); for ch in literal.chars() { if ch == '/' || ch == '\\' { if !segment.is_empty() { append_path_segment(builder, &segment); segment.clear(); } builder.append_o(&ch.to_string()); } else { segment.push(ch); } } if !segment.is_empty() { append_path_segment(builder, &segment); } } fn append_path_segment(builder: &mut CharBuilder, segment: &str) { if segment.is_empty() { return; } if is_path_season_segment(segment) { builder.append_entity(segment, Entity::PathSeason); } else if is_path_tag_segment(segment) { builder.append_entity(segment, Entity::Tag); } else { builder.append_o(segment); } } fn is_path_season_segment(segment: &str) -> bool { static RE: OnceLock = OnceLock::new(); let trimmed = segment.trim(); let re = RE.get_or_init(|| { Regex::new(r"(?i)^(?:season\s*0?\d{1,2}|s0?\d{1,2}|0?[1-9]|1[0-9]|第[一二三四五六七八九十\d]+[季期部])$") .expect("path season regex compiles") }); re.is_match(trimmed) } fn is_path_tag_segment(segment: &str) -> bool { static YEAR_RE: OnceLock = OnceLock::new(); let trimmed = segment.trim(); if matches!( trimmed.to_ascii_lowercase().as_str(), "gekijouban" | "movie" | "movies" | "anime movies" | "films" | "bdmovie" | "tv" | "extras" | "sps" | "specials" | "ova" | "oad" | "ncop" | "nced" | "pv" | "cm" | "trailer" | "menu" ) || matches!(trimmed, "劇場版") { return true; } YEAR_RE .get_or_init(|| Regex::new(r"^(?:19|20)\d{2}$").expect("year regex compiles")) .is_match(trimmed) } #[allow(dead_code)] fn entities_for_text<'a>(record: &'a Record, needle: &str) -> Vec<&'a str> { let Some(start) = record.filename.find(needle) else { return Vec::new(); }; let end = start + needle.len(); char_byte_offsets(&record.filename) .iter() .enumerate() .filter(|(_, byte_idx)| **byte_idx >= start && **byte_idx < end) .map(|(idx, _)| record.labels[idx].as_str()) .collect() } #[allow(dead_code)] fn fixture_record(filename: &str, spans: &[(&str, Entity)]) -> Record { let mut byte_spans = Vec::new(); for (needle, entity) in spans { let start = filename .find(needle) .unwrap_or_else(|| panic!("fixture missing span {needle} in {filename}")); byte_spans.push((start, start + needle.len(), *entity)); } char_record_from_spans( filename, &byte_spans, "test", "test", "test", "numeric_title", ) } #[cfg(test)] mod tests { use super::*; fn assert_all_entity(record: &Record, needle: &str, entity: &str) { let labels = entities_for_text(record, needle); assert!(!labels.is_empty(), "missing text span {needle}"); assert!( labels.iter().all(|label| label.ends_with(entity)), "{needle} labels were {labels:?}, expected {entity}" ); } fn assert_no_entity(record: &Record, needle: &str, entity: &str) { let labels = entities_for_text(record, needle); assert!(!labels.is_empty(), "missing text span {needle}"); assert!( labels.iter().all(|label| !label.ends_with(entity)), "{needle} labels unexpectedly included {entity}: {labels:?}" ); } fn all_test_labels() -> HashSet { let mut labels = HashSet::from(["O".to_string()]); for entity in [ Entity::TitleLatin, Entity::PathTitleLatin, Entity::PathSeason, Entity::Season, Entity::Episode, Entity::Special, Entity::Group, Entity::Resolution, Entity::Source, Entity::Tag, ] { labels.insert(entity.b_label().to_string()); labels.insert(entity.i_label().to_string()); } labels } #[test] fn numeric_title_7_nin_nced_keeps_number_in_title() { let record = fixture_record( "Yamada-kun to 7-nin no Majo [NCED]", &[ ("Yamada-kun to 7-nin no Majo", Entity::TitleLatin), ("NCED", Entity::Special), ], ); assert_all_entity(&record, "7-nin", "TITLE_LATIN"); assert_all_entity(&record, "NCED", "SPECIAL"); assert!(!record.labels.iter().any(|label| label.ends_with("EPISODE"))); } #[test] fn numeric_title_91_days_ncop2_keeps_number_in_title() { let record = fixture_record( "91 Days [NCOP2]", &[("91 Days", Entity::TitleLatin), ("NCOP2", Entity::Special)], ); assert_all_entity(&record, "91 Days", "TITLE_LATIN"); assert_all_entity(&record, "NCOP2", "SPECIAL"); assert_no_entity(&record, "91", "EPISODE"); } #[test] fn special_number_after_pv_is_special_not_episode() { let recipe = Recipe { template_id: "tpl_test".to_string(), template: "TEXT SEP SEASON SEP BRACKET_SPECIAL BRACKET_EPISODE".to_string(), roles: vec![ "TITLE".to_string(), "O".to_string(), "SEASON".to_string(), "O".to_string(), "SPECIAL".to_string(), "EPISODE".to_string(), ], title_spans: Vec::new(), confidence: Some("high".to_string()), _count: Some(1), }; let record = build_numeric_record(&recipe, "100-nin no Kanojo", 2).unwrap(); assert_all_entity(&record, "100-nin", "TITLE_LATIN"); assert_all_entity(&record, "S01", "SEASON"); assert_all_entity(&record, "PV", "SPECIAL"); let pv_start = record.filename.find("PV").unwrap(); let special_number_labels = entities_for_text_after(&record, "01", pv_start); assert!( special_number_labels .iter() .all(|label| label.ends_with("SPECIAL")), "special number labels were {special_number_labels:?}" ); } #[test] fn numeric_generation_drops_multi_title_templates() { let recipe = Recipe { template_id: "tpl_multi_title".to_string(), template: "TEXT SEP TEXT SEP EPISODE".to_string(), roles: vec![ "TITLE".to_string(), "O".to_string(), "TITLE".to_string(), "O".to_string(), "EPISODE".to_string(), ], title_spans: Vec::new(), confidence: Some("high".to_string()), _count: Some(1), }; assert!(build_numeric_record(&recipe, "91 Days", 0).is_none()); } #[test] fn numeric_generation_splits_merged_logical_title_span() { let recipe = Recipe { template_id: "tpl_merged_title".to_string(), template: "TEXT SEP TEXT SEP TEXT SEP EPISODE".to_string(), roles: vec![ "TITLE".to_string(), "O".to_string(), "TITLE".to_string(), "O".to_string(), "TITLE".to_string(), "O".to_string(), "EPISODE".to_string(), ], title_spans: vec![TitleSpan { role_indices: vec![0, 2, 4], }], confidence: Some("high".to_string()), _count: Some(1), }; let record = build_numeric_record(&recipe, "500 Days Until the Dungeon Closes", 0).unwrap(); assert!(record.filename.contains("500 Days")); assert!(record.filename.contains("Until the")); assert!(record.filename.contains("Dungeon Closes")); assert_eq!( record .filename .matches("500 Days Until the Dungeon Closes") .count(), 1 ); assert_all_entity(&record, "500 Days", "TITLE_LATIN"); assert_all_entity(&record, "Until the", "TITLE_LATIN"); assert_all_entity(&record, "Dungeon Closes", "TITLE_LATIN"); } #[test] fn path_title_season_episode_labels_are_projected() { let base = char_record_from_spans( "03.mkv", &[(0, 2, Entity::Episode)], "test", "leaf", "leaf", "path_aug", ); let record = build_path_record( "/mnt/media/anime/{title}/Season 01/{filename}", PathSeedKind::Series, "Title", "Season 01", &base, ) .unwrap(); assert_eq!(record.filename, "/mnt/media/anime/Title/Season 01/03.mkv"); assert_all_entity(&record, "Title", "PATH_TITLE_LATIN"); assert_all_entity(&record, "Season 01", "PATH_SEASON"); assert_all_entity(&record, "03", "EPISODE"); } #[test] fn confusing_path_dirs_are_tags_not_path_season() { let base = char_record_from_spans("file.mkv", &[], "test", "leaf", "leaf", "path_aug"); let record = build_path_record( "/mnt/media/anime/{title}/Gekijouban/2004/{filename}", PathSeedKind::Confuser, "Naruto", "Season 01", &base, ) .unwrap(); assert_eq!( record.filename, "/mnt/media/anime/Naruto/Gekijouban/2004/file.mkv" ); assert_all_entity(&record, "Naruto", "PATH_TITLE_LATIN"); assert_all_entity(&record, "Gekijouban", "TAG"); assert_all_entity(&record, "2004", "TAG"); assert_no_entity(&record, "2004", "PATH_SEASON"); } #[test] fn movie_path_rejects_series_episode_base_and_accepts_movie_base() { let series_base = char_record_from_spans( "03.mkv", &[(0, 2, Entity::Episode)], "test", "leaf", "leaf", "path_aug", ); assert!(build_path_record( "Movies/{title}/{filename}", PathSeedKind::Movie, "Area 88", "Season 01", &series_base, ) .is_none()); let movie_base = char_record_from_spans( "Movie.mkv", &[(0, 5, Entity::Tag)], "test", "leaf_movie", "leaf_movie", "path_aug", ); let record = build_path_record( "Movies/{title}/{filename}", PathSeedKind::Movie, "Area 88", "Season 01", &movie_base, ) .unwrap(); assert_eq!(record.synthetic_kind, "path_movie"); assert_all_entity(&record, "Area 88", "PATH_TITLE_LATIN"); assert_all_entity(&record, "Movie", "TAG"); assert!(!record.filename.contains("Season 01")); } #[test] fn series_path_rejects_movie_base() { let movie_base = char_record_from_spans( "Gekijouban.mkv", &[(0, 10, Entity::Tag)], "test", "leaf_movie", "leaf_movie", "path_aug", ); assert!(build_path_record( "Anime/{title}/Season 01/{filename}", PathSeedKind::Series, "No.6", "Season 01", &movie_base, ) .is_none()); } #[test] fn typed_path_seed_parser_distinguishes_movie_special_and_confuser() { assert_eq!( parse_path_seed_line("movie|Movies/{title}/{filename}") .unwrap() .kind, PathSeedKind::Movie ); assert_eq!( parse_path_seed_line("special|SPs/{title}/{filename}") .unwrap() .kind, PathSeedKind::Special ); assert_eq!( parse_path_seed_line("# confuser|Bangumi/{title}/2004/{filename}") .unwrap() .kind, PathSeedKind::Confuser ); } #[test] fn path_kind_stats_increment_separate_manifest_buckets() { let mut stats = BuildStats::default(); let mut records = Vec::new(); let mut seen = HashSet::new(); let label_set = all_test_labels(); let movie = char_record_from_spans( "Movies/Area 88/Movie.mkv", &[(7, 14, Entity::PathTitleLatin), (15, 20, Entity::Tag)], "test", "movie", "movie", "path_movie", ); push_record(movie, &mut records, &mut seen, &mut stats, &label_set).unwrap(); assert_eq!(stats.path_rows, 1); assert_eq!(stats.path_movie_rows, 1); assert_eq!(stats.path_series_rows, 0); assert_eq!(stats.path_special_rows, 0); } #[test] fn rust_record_validation_rejects_orphan_i_tag_and_missing_title() { let label_set = all_test_labels(); let mut orphan = fixture_record("No.6.mkv", &[("No.6", Entity::TitleLatin)]); orphan.labels[0] = "I-TITLE_LATIN".to_string(); assert!(validate_record_labels(&orphan, &label_set) .unwrap_err() .to_string() .contains("orphan I-tag")); let missing_title = char_record_from_spans( "03.mkv", &[(0, 2, Entity::Episode)], "test", "episode", "episode", "path_series", ); assert!(validate_record_labels(&missing_title, &label_set) .unwrap_err() .to_string() .contains("missing title span")); } #[test] fn manifest_rows_round_trip_json_shape() { let record = fixture_record( "86 [01][1080p]", &[ ("86", Entity::TitleLatin), ("01", Entity::Episode), ("1080p", Entity::Resolution), ], ); let encoded = serde_json::to_value(&record).unwrap(); assert_eq!( encoded["filename"], Value::String("86 [01][1080p]".to_string()) ); assert_eq!( encoded["tokenizer_variant"], Value::String("char".to_string()) ); } } #[cfg(test)] fn entities_for_text_after<'a>(record: &'a Record, needle: &str, start_at: usize) -> Vec<&'a str> { let Some(start) = record.filename[start_at..].find(needle) else { return Vec::new(); }; let start = start + start_at; let end = start + needle.len(); char_byte_offsets(&record.filename) .iter() .enumerate() .filter(|(_, byte_idx)| **byte_idx >= start && **byte_idx < end) .map(|(idx, _)| record.labels[idx].as_str()) .collect() }