Token Classification
Transformers
ONNX
Safetensors
English
Japanese
Chinese
bert
anime
filename-parsing
Eval Results (legacy)
Instructions to use ModerRAS/AniFileBERT with libraries, inference providers, notebooks, and local apps. Follow these links to get started.
- Libraries
- Transformers
How to use ModerRAS/AniFileBERT with Transformers:
# Use a pipeline as a high-level helper from transformers import pipeline pipe = pipeline("token-classification", model="ModerRAS/AniFileBERT")# Load model directly from transformers import AutoTokenizer, AutoModelForTokenClassification tokenizer = AutoTokenizer.from_pretrained("ModerRAS/AniFileBERT") model = AutoModelForTokenClassification.from_pretrained("ModerRAS/AniFileBERT") - Notebooks
- Google Colab
- Kaggle
| use std::collections::{BTreeMap, BTreeSet}; | |
| use std::fs::{self, File}; | |
| use std::io::{BufRead, BufReader, BufWriter}; | |
| use std::path::{Path, PathBuf}; | |
| use std::sync::LazyLock; | |
| use anyhow::{Context, Result}; | |
| use clap::Parser; | |
| use rayon::prelude::*; | |
| use regex::Regex; | |
| use serde::{Deserialize, Serialize}; | |
| const RADIX_VERSION: &str = "prefix-radix-v1"; | |
| const DAG_VERSION: &str = "prefix-dag-v1"; | |
| static EPISODE_PATTERNS: LazyLock<Vec<Regex>> = LazyLock::new(|| { | |
| vec![ | |
| Regex::new(r"(?i)\bS\d{1,2}E\d{1,4}\b").unwrap(), | |
| Regex::new(r"(?i)\bEP?\.?\s*\d{1,4}\b").unwrap(), | |
| Regex::new(r"第\s*\d{1,4}\s*[話话集回]").unwrap(), | |
| Regex::new(r"\b\d{1,4}\s*[話话集回]").unwrap(), | |
| Regex::new(r"[\[(]\s*\d{1,4}\s*[\])]").unwrap(), | |
| ] | |
| }); | |
| static DIGITS: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"\d+").unwrap()); | |
| struct Args { | |
| input: PathBuf, | |
| output: PathBuf, | |
| min_count: usize, | |
| example_count: usize, | |
| from_prefix_graph: Option<PathBuf>, | |
| dag_output: Option<PathBuf>, | |
| } | |
| struct InputRecord { | |
| value: String, | |
| uses_path: bool, | |
| has_trailing_hash: bool, | |
| _has_digits: bool, | |
| _digit_skeleton: String, | |
| count: usize, | |
| } | |
| struct PatternObservation { | |
| source_index: usize, | |
| prefix: String, | |
| digit_skeleton: String, | |
| suffix: String, | |
| value: String, | |
| uses_path: bool, | |
| has_trailing_hash: bool, | |
| count: usize, | |
| } | |
| struct GroupBuilder { | |
| prefix: String, | |
| digit_skeleton: String, | |
| count: usize, | |
| weight: usize, | |
| uses_path_count: usize, | |
| has_trailing_hash_count: usize, | |
| suffix_examples: Vec<String>, | |
| value_examples: Vec<String>, | |
| } | |
| struct TrieNode { | |
| id: usize, | |
| edge_label: String, | |
| depth: usize, | |
| parent: Option<usize>, | |
| children_by_label: BTreeMap<String, usize>, | |
| subtree_patterns: usize, | |
| subtree_weight: usize, | |
| } | |
| struct GraphOutput { | |
| meta: Meta, | |
| nodes: Vec<OutputNode>, | |
| terminals: Vec<OutputTerminal>, | |
| } | |
| struct Meta { | |
| input: String, | |
| output: String, | |
| input_records: usize, | |
| observations: usize, | |
| no_episode_prefix: usize, | |
| groups: usize, | |
| nodes: usize, | |
| max_depth: usize, | |
| grouped_weight: usize, | |
| min_count: usize, | |
| example_count: usize, | |
| tokenizer: TokenizerMeta, | |
| } | |
| struct TokenizerMeta { | |
| version: &'static str, | |
| notes: Vec<&'static str>, | |
| } | |
| struct OutputNode { | |
| id: usize, | |
| edge_label: String, | |
| depth: usize, | |
| children: Vec<usize>, | |
| subtree_patterns: usize, | |
| subtree_weight: usize, | |
| } | |
| struct OutputTerminal { | |
| node_id: usize, | |
| prefix: String, | |
| digit_skeleton: String, | |
| count: usize, | |
| weight: usize, | |
| uses_path_count: usize, | |
| has_trailing_hash_count: usize, | |
| suffix_examples: Vec<String>, | |
| value_examples: Vec<String>, | |
| annotations: TerminalAnnotations, | |
| } | |
| struct TerminalAnnotations { | |
| episode_title_suffixes: Vec<String>, | |
| media_suffixes: Vec<String>, | |
| title_candidates: Vec<String>, | |
| needs_llm_review: bool, | |
| llm_label: Option<String>, | |
| notes: Option<String>, | |
| } | |
| struct SourceGraph { | |
| meta: serde_json::Value, | |
| nodes: Vec<SourceNode>, | |
| terminals: Vec<SourceTerminal>, | |
| } | |
| struct SourceNode { | |
| id: usize, | |
| edge_label: String, | |
| depth: usize, | |
| children: Vec<usize>, | |
| } | |
| struct SourceTerminal { | |
| terminal_id: Option<String>, | |
| node_id: usize, | |
| prefix: String, | |
| digit_skeleton: String, | |
| count: usize, | |
| weight: usize, | |
| uses_path_count: usize, | |
| has_trailing_hash_count: usize, | |
| suffix_examples: Vec<String>, | |
| value_examples: Vec<String>, | |
| annotations: serde_json::Value, | |
| } | |
| struct DagOutput { | |
| meta: DagMeta, | |
| root: usize, | |
| nodes: Vec<DagNode>, | |
| terminals: Vec<DagTerminal>, | |
| } | |
| struct DagMeta { | |
| version: &'static str, | |
| input: String, | |
| output: String, | |
| source_nodes: usize, | |
| source_terminals: usize, | |
| dag_nodes: usize, | |
| dag_edges: usize, | |
| root: usize, | |
| max_depth: usize, | |
| merged_nodes: usize, | |
| preserves_digits: bool, | |
| merge_strategy: &'static str, | |
| notes: Vec<&'static str>, | |
| } | |
| struct DagNode { | |
| id: usize, | |
| terminal: bool, | |
| children: Vec<DagEdge>, | |
| incoming_count: usize, | |
| reachable_terminals: usize, | |
| reachable_weight: usize, | |
| } | |
| struct DagEdge { | |
| label: String, | |
| target: usize, | |
| } | |
| struct DagTerminal { | |
| terminal_id: String, | |
| node_id: usize, | |
| prefix: String, | |
| digit_skeleton: String, | |
| count: usize, | |
| weight: usize, | |
| uses_path_count: usize, | |
| has_trailing_hash_count: usize, | |
| suffix_examples: Vec<String>, | |
| value_examples: Vec<String>, | |
| annotations: serde_json::Value, | |
| } | |
| struct NodeSignature { | |
| terminal: bool, | |
| children: Vec<(String, usize)>, | |
| } | |
| struct TempDagNode { | |
| terminal: bool, | |
| children: Vec<DagEdge>, | |
| } | |
| fn main() -> Result<()> { | |
| let args = Args::parse(); | |
| if args.from_prefix_graph.is_some() || args.dag_output.is_some() { | |
| let dag = build_dag_from_args(&args)?; | |
| println!("{}", serde_json::to_string_pretty(&dag.meta)?); | |
| } else { | |
| let graph = build_graph(&args)?; | |
| println!("{}", serde_json::to_string_pretty(&graph.meta)?); | |
| } | |
| Ok(()) | |
| } | |
| fn build_graph(args: &Args) -> Result<GraphOutput> { | |
| ensure_output_parent(&args.output)?; | |
| let records = read_records(&args.input)?; | |
| let input_records = records.len(); | |
| let mut observations = records | |
| .par_iter() | |
| .enumerate() | |
| .filter_map(|(source_index, record)| build_observation(source_index, record)) | |
| .collect::<Vec<_>>(); | |
| observations.sort_by(|left, right| { | |
| left.prefix | |
| .cmp(&right.prefix) | |
| .then(left.source_index.cmp(&right.source_index)) | |
| }); | |
| let no_episode_prefix = input_records.saturating_sub(observations.len()); | |
| let groups = aggregate_observations(&observations, args.min_count, args.example_count); | |
| let (nodes, terminals, max_depth, grouped_weight) = build_trie(&groups); | |
| let graph = GraphOutput { | |
| meta: Meta { | |
| input: display_path(&args.input), | |
| output: display_path(&args.output), | |
| input_records, | |
| observations: observations.len(), | |
| no_episode_prefix, | |
| groups: terminals.len(), | |
| nodes: nodes.len(), | |
| max_depth, | |
| grouped_weight, | |
| min_count: args.min_count, | |
| example_count: args.example_count, | |
| tokenizer: TokenizerMeta { | |
| version: RADIX_VERSION, | |
| notes: vec![ | |
| "Episode prefixes are detected with the legacy regex/boundary rules.", | |
| "Graph insertion preserves original prefix digits; digit_skeleton is secondary metadata only.", | |
| "Graph nodes form a compressed radix trie over complete original prefix strings.", | |
| "Edge labels split only at actual branch points; punctuation and separators do not force levels.", | |
| ], | |
| }, | |
| }, | |
| nodes, | |
| terminals, | |
| }; | |
| let output = File::create(&args.output) | |
| .with_context(|| format!("failed to create {}", args.output.display()))?; | |
| let writer = BufWriter::new(output); | |
| serde_json::to_writer_pretty(writer, &graph) | |
| .with_context(|| format!("failed to write {}", args.output.display()))?; | |
| Ok(graph) | |
| } | |
| fn build_dag_from_args(args: &Args) -> Result<DagOutput> { | |
| let input = args | |
| .from_prefix_graph | |
| .as_ref() | |
| .unwrap_or(&args.output) | |
| .to_path_buf(); | |
| let output = args | |
| .dag_output | |
| .as_ref() | |
| .cloned() | |
| .unwrap_or_else(|| PathBuf::from("datasets\\AnimeName\\dmhy_prefix_dag.json")); | |
| build_dag(&input, &output) | |
| } | |
| fn build_dag(input_path: &Path, output_path: &Path) -> Result<DagOutput> { | |
| ensure_output_parent(output_path)?; | |
| let input = File::open(input_path) | |
| .with_context(|| format!("failed to open {}", input_path.display()))?; | |
| let reader = BufReader::new(input); | |
| let source: SourceGraph = serde_json::from_reader(reader) | |
| .with_context(|| format!("failed to parse {}", input_path.display()))?; | |
| validate_source_graph(&source)?; | |
| let dag = minimize_source_graph(&source, input_path, output_path)?; | |
| let output = File::create(output_path) | |
| .with_context(|| format!("failed to create {}", output_path.display()))?; | |
| let writer = BufWriter::new(output); | |
| serde_json::to_writer_pretty(writer, &dag) | |
| .with_context(|| format!("failed to write {}", output_path.display()))?; | |
| Ok(dag) | |
| } | |
| fn validate_source_graph(source: &SourceGraph) -> Result<()> { | |
| for (index, node) in source.nodes.iter().enumerate() { | |
| anyhow::ensure!( | |
| node.id == index, | |
| "source node id {} appears at index {}", | |
| node.id, | |
| index | |
| ); | |
| for &child_id in &node.children { | |
| anyhow::ensure!( | |
| child_id < source.nodes.len(), | |
| "source node {} references missing child {}", | |
| node.id, | |
| child_id | |
| ); | |
| } | |
| } | |
| for terminal in &source.terminals { | |
| anyhow::ensure!( | |
| terminal.node_id < source.nodes.len(), | |
| "terminal prefix {:?} references missing node {}", | |
| terminal.prefix, | |
| terminal.node_id | |
| ); | |
| } | |
| Ok(()) | |
| } | |
| fn minimize_source_graph( | |
| source: &SourceGraph, | |
| input_path: &Path, | |
| output_path: &Path, | |
| ) -> Result<DagOutput> { | |
| let terminal_source_nodes = source | |
| .terminals | |
| .iter() | |
| .map(|terminal| terminal.node_id) | |
| .collect::<BTreeSet<_>>(); | |
| let postorder = source_postorder(source); | |
| let mut source_to_temp_dag = vec![usize::MAX; source.nodes.len()]; | |
| let mut signatures = BTreeMap::<NodeSignature, usize>::new(); | |
| let mut temp_nodes = Vec::<TempDagNode>::new(); | |
| for source_id in postorder { | |
| let source_node = &source.nodes[source_id]; | |
| let mut children = source_node | |
| .children | |
| .iter() | |
| .map(|&child_id| { | |
| let child = &source.nodes[child_id]; | |
| let target = source_to_temp_dag[child_id]; | |
| anyhow::ensure!( | |
| target != usize::MAX, | |
| "source child {} was not canonicalized before parent {}", | |
| child_id, | |
| source_id | |
| ); | |
| Ok((child.edge_label.clone(), target)) | |
| }) | |
| .collect::<Result<Vec<_>>>()?; | |
| children.sort(); | |
| let signature = NodeSignature { | |
| terminal: terminal_source_nodes.contains(&source_id), | |
| children, | |
| }; | |
| let temp_id = if let Some(&existing) = signatures.get(&signature) { | |
| existing | |
| } else { | |
| let temp_id = temp_nodes.len(); | |
| temp_nodes.push(TempDagNode { | |
| terminal: signature.terminal, | |
| children: signature | |
| .children | |
| .iter() | |
| .map(|(label, target)| DagEdge { | |
| label: label.clone(), | |
| target: *target, | |
| }) | |
| .collect(), | |
| }); | |
| signatures.insert(signature, temp_id); | |
| temp_id | |
| }; | |
| source_to_temp_dag[source_id] = temp_id; | |
| } | |
| let root_temp_id = source_to_temp_dag[0]; | |
| let (temp_to_dag, mut dag_nodes) = renumber_dag(root_temp_id, &temp_nodes); | |
| for node in &mut dag_nodes { | |
| for edge in &mut node.children { | |
| edge.target = temp_to_dag[edge.target]; | |
| } | |
| } | |
| let terminals = source | |
| .terminals | |
| .iter() | |
| .enumerate() | |
| .map(|(index, terminal)| { | |
| let temp_id = source_to_temp_dag[terminal.node_id]; | |
| DagTerminal { | |
| terminal_id: terminal | |
| .terminal_id | |
| .clone() | |
| .unwrap_or_else(|| format!("t{}", index)), | |
| node_id: temp_to_dag[temp_id], | |
| prefix: terminal.prefix.clone(), | |
| digit_skeleton: terminal.digit_skeleton.clone(), | |
| count: terminal.count, | |
| weight: terminal.weight, | |
| uses_path_count: terminal.uses_path_count, | |
| has_trailing_hash_count: terminal.has_trailing_hash_count, | |
| suffix_examples: terminal.suffix_examples.clone(), | |
| value_examples: terminal.value_examples.clone(), | |
| annotations: terminal.annotations.clone(), | |
| } | |
| }) | |
| .collect::<Vec<_>>(); | |
| fill_dag_counts(&source, &source_to_temp_dag, &temp_to_dag, &mut dag_nodes); | |
| let dag_edges = dag_nodes.iter().map(|node| node.children.len()).sum(); | |
| let max_depth = dag_max_depth(&dag_nodes); | |
| Ok(DagOutput { | |
| meta: DagMeta { | |
| version: DAG_VERSION, | |
| input: display_path(input_path), | |
| output: display_path(output_path), | |
| source_nodes: source.nodes.len(), | |
| source_terminals: source.terminals.len(), | |
| dag_nodes: dag_nodes.len(), | |
| dag_edges, | |
| root: 0, | |
| max_depth, | |
| merged_nodes: source.nodes.len().saturating_sub(dag_nodes.len()), | |
| preserves_digits: true, | |
| merge_strategy: | |
| "bottom-up suffix equivalence over radix trie nodes; edge labels remain raw substrings", | |
| notes: vec![ | |
| "Terminal prefixes and digit_skeleton values are copied from the source graph; digits are not normalized for merging.", | |
| "Node reachable_terminals and reachable_weight count terminal instances mapped onto the shared DAG suffix, so shared nodes report aggregate suffix usage rather than one caller-specific path.", | |
| ], | |
| }, | |
| root: 0, | |
| nodes: dag_nodes, | |
| terminals, | |
| }) | |
| } | |
| fn source_postorder(source: &SourceGraph) -> Vec<usize> { | |
| let mut order = Vec::with_capacity(source.nodes.len()); | |
| let mut stack = vec![(0, false)]; | |
| while let Some((node_id, expanded)) = stack.pop() { | |
| if expanded { | |
| order.push(node_id); | |
| } else { | |
| stack.push((node_id, true)); | |
| for &child_id in source.nodes[node_id].children.iter().rev() { | |
| stack.push((child_id, false)); | |
| } | |
| } | |
| } | |
| order | |
| } | |
| fn renumber_dag(root_temp_id: usize, temp_nodes: &[TempDagNode]) -> (Vec<usize>, Vec<DagNode>) { | |
| let mut temp_to_dag = vec![usize::MAX; temp_nodes.len()]; | |
| let mut dag_nodes = Vec::<DagNode>::new(); | |
| let mut stack = vec![root_temp_id]; | |
| while let Some(temp_id) = stack.pop() { | |
| if temp_to_dag[temp_id] != usize::MAX { | |
| continue; | |
| } | |
| let dag_id = dag_nodes.len(); | |
| temp_to_dag[temp_id] = dag_id; | |
| dag_nodes.push(DagNode { | |
| id: dag_id, | |
| terminal: temp_nodes[temp_id].terminal, | |
| children: temp_nodes[temp_id].children.clone(), | |
| incoming_count: 0, | |
| reachable_terminals: 0, | |
| reachable_weight: 0, | |
| }); | |
| for edge in temp_nodes[temp_id].children.iter().rev() { | |
| stack.push(edge.target); | |
| } | |
| } | |
| (temp_to_dag, dag_nodes) | |
| } | |
| fn fill_dag_counts( | |
| source: &SourceGraph, | |
| source_to_temp_dag: &[usize], | |
| temp_to_dag: &[usize], | |
| nodes: &mut [DagNode], | |
| ) { | |
| let mut terminal_ids_by_node = vec![BTreeSet::<usize>::new(); nodes.len()]; | |
| let all_edges = nodes | |
| .iter() | |
| .flat_map(|node| node.children.iter().map(|edge| edge.target)) | |
| .collect::<Vec<_>>(); | |
| for target in all_edges { | |
| nodes[target].incoming_count += 1; | |
| } | |
| let source_parents = source_parent_map(source); | |
| for (terminal_index, terminal) in source.terminals.iter().enumerate() { | |
| let mut source_node_id = Some(terminal.node_id); | |
| while let Some(node_id) = source_node_id { | |
| let temp_id = source_to_temp_dag[node_id]; | |
| let dag_id = temp_to_dag[temp_id]; | |
| terminal_ids_by_node[dag_id].insert(terminal_index); | |
| source_node_id = source_parents[node_id]; | |
| } | |
| } | |
| for (node_id, terminal_ids) in terminal_ids_by_node.into_iter().enumerate() { | |
| nodes[node_id].reachable_terminals = terminal_ids.len(); | |
| nodes[node_id].reachable_weight = terminal_ids | |
| .into_iter() | |
| .map(|terminal_id| source.terminals[terminal_id].weight) | |
| .sum(); | |
| } | |
| } | |
| fn source_parent_map(source: &SourceGraph) -> Vec<Option<usize>> { | |
| let mut parents = vec![None; source.nodes.len()]; | |
| for node in &source.nodes { | |
| for &child_id in &node.children { | |
| parents[child_id] = Some(node.id); | |
| } | |
| } | |
| parents | |
| } | |
| fn dag_max_depth(nodes: &[DagNode]) -> usize { | |
| let mut max_depth = 0; | |
| let mut stack = vec![(0, 0)]; | |
| while let Some((node_id, depth)) = stack.pop() { | |
| max_depth = max_depth.max(depth); | |
| for edge in &nodes[node_id].children { | |
| stack.push((edge.target, depth + 1)); | |
| } | |
| } | |
| max_depth | |
| } | |
| fn ensure_output_parent(output: &Path) -> Result<()> { | |
| if let Some(parent) = output.parent() { | |
| if !parent.as_os_str().is_empty() { | |
| fs::create_dir_all(parent) | |
| .with_context(|| format!("failed to create {}", parent.display()))?; | |
| } | |
| } | |
| Ok(()) | |
| } | |
| fn read_records(input_path: &Path) -> Result<Vec<InputRecord>> { | |
| let input = File::open(input_path) | |
| .with_context(|| format!("failed to open {}", input_path.display()))?; | |
| let reader = BufReader::new(input); | |
| let mut records = Vec::new(); | |
| for (line_number, line) in reader.lines().enumerate() { | |
| let line = line.with_context(|| { | |
| format!( | |
| "failed to read line {} from {}", | |
| line_number + 1, | |
| input_path.display() | |
| ) | |
| })?; | |
| let line = line.trim(); | |
| if line.is_empty() { | |
| continue; | |
| } | |
| let record = serde_json::from_str(line).with_context(|| { | |
| format!( | |
| "failed to parse line {} from {}", | |
| line_number + 1, | |
| input_path.display() | |
| ) | |
| })?; | |
| records.push(record); | |
| } | |
| Ok(records) | |
| } | |
| fn build_observation(source_index: usize, record: &InputRecord) -> Option<PatternObservation> { | |
| let (prefix, suffix) = find_episode_prefix(&record.value)?; | |
| if prefix.is_empty() { | |
| return None; | |
| } | |
| let digit_skeleton = digit_skeleton(&prefix); | |
| Some(PatternObservation { | |
| source_index, | |
| prefix, | |
| digit_skeleton, | |
| suffix, | |
| value: record.value.clone(), | |
| uses_path: record.uses_path, | |
| has_trailing_hash: record.has_trailing_hash, | |
| count: record.count, | |
| }) | |
| } | |
| fn aggregate_observations( | |
| observations: &[PatternObservation], | |
| min_count: usize, | |
| example_count: usize, | |
| ) -> Vec<GroupBuilder> { | |
| let mut groups = BTreeMap::<String, GroupBuilder>::new(); | |
| for observation in observations { | |
| let group = groups | |
| .entry(observation.prefix.clone()) | |
| .or_insert_with(|| GroupBuilder { | |
| prefix: observation.prefix.clone(), | |
| digit_skeleton: observation.digit_skeleton.clone(), | |
| ..GroupBuilder::default() | |
| }); | |
| group.count += 1; | |
| group.weight += observation.count; | |
| group.uses_path_count += observation.count * usize::from(observation.uses_path); | |
| group.has_trailing_hash_count += | |
| observation.count * usize::from(observation.has_trailing_hash); | |
| if !observation.suffix.is_empty() && group.suffix_examples.len() < example_count { | |
| group.suffix_examples.push(observation.suffix.clone()); | |
| } | |
| if group.value_examples.len() < example_count { | |
| group.value_examples.push(observation.value.clone()); | |
| } | |
| } | |
| groups | |
| .into_values() | |
| .filter(|group| group.count >= min_count) | |
| .collect() | |
| } | |
| fn build_trie(groups: &[GroupBuilder]) -> (Vec<OutputNode>, Vec<OutputTerminal>, usize, usize) { | |
| let mut nodes = vec![TrieNode::root()]; | |
| let mut terminal_node_ids = Vec::with_capacity(groups.len()); | |
| for group in groups { | |
| let (terminal_node_id, path) = insert_prefix(&mut nodes, &group.prefix); | |
| terminal_node_ids.push(terminal_node_id); | |
| for id in path { | |
| nodes[id].subtree_patterns += group.count; | |
| nodes[id].subtree_weight += group.weight; | |
| } | |
| } | |
| assign_depths(&mut nodes); | |
| let terminals = groups | |
| .iter() | |
| .zip(terminal_node_ids) | |
| .map(|(group, node_id)| OutputTerminal { | |
| node_id, | |
| prefix: group.prefix.clone(), | |
| digit_skeleton: group.digit_skeleton.clone(), | |
| count: group.count, | |
| weight: group.weight, | |
| uses_path_count: group.uses_path_count, | |
| has_trailing_hash_count: group.has_trailing_hash_count, | |
| suffix_examples: group.suffix_examples.clone(), | |
| value_examples: group.value_examples.clone(), | |
| annotations: TerminalAnnotations::default(), | |
| }) | |
| .collect::<Vec<_>>(); | |
| let max_depth = nodes.iter().map(|node| node.depth).max().unwrap_or(0); | |
| let grouped_weight = groups.iter().map(|group| group.weight).sum(); | |
| let output_nodes = nodes | |
| .into_iter() | |
| .map(|node| OutputNode { | |
| id: node.id, | |
| edge_label: node.edge_label, | |
| depth: node.depth, | |
| children: node.children_by_label.into_values().collect(), | |
| subtree_patterns: node.subtree_patterns, | |
| subtree_weight: node.subtree_weight, | |
| }) | |
| .collect(); | |
| (output_nodes, terminals, max_depth, grouped_weight) | |
| } | |
| impl TrieNode { | |
| fn root() -> Self { | |
| Self { | |
| id: 0, | |
| edge_label: String::new(), | |
| depth: 0, | |
| parent: None, | |
| children_by_label: BTreeMap::new(), | |
| subtree_patterns: 0, | |
| subtree_weight: 0, | |
| } | |
| } | |
| } | |
| fn insert_prefix(nodes: &mut Vec<TrieNode>, prefix: &str) -> (usize, Vec<usize>) { | |
| let mut current = 0; | |
| let mut remaining = prefix; | |
| let mut path = vec![current]; | |
| loop { | |
| if remaining.is_empty() { | |
| return (current, path); | |
| } | |
| let matching_child = | |
| nodes[current] | |
| .children_by_label | |
| .iter() | |
| .find_map(|(label, &child_id)| { | |
| let common_len = common_prefix_len(remaining, label); | |
| (common_len > 0).then(|| (label.clone(), child_id, common_len)) | |
| }); | |
| let Some((child_label, child_id, common_len)) = matching_child else { | |
| let child_id = push_node(nodes, current, remaining.to_owned()); | |
| nodes[current] | |
| .children_by_label | |
| .insert(remaining.to_owned(), child_id); | |
| path.push(child_id); | |
| return (child_id, path); | |
| }; | |
| if common_len == child_label.len() { | |
| current = child_id; | |
| remaining = &remaining[common_len..]; | |
| path.push(current); | |
| continue; | |
| } | |
| let shared_label = child_label[..common_len].to_owned(); | |
| let old_suffix = child_label[common_len..].to_owned(); | |
| let new_suffix = remaining[common_len..].to_owned(); | |
| let split_id = push_node(nodes, current, shared_label.clone()); | |
| nodes[current].children_by_label.remove(&child_label); | |
| nodes[current] | |
| .children_by_label | |
| .insert(shared_label, split_id); | |
| nodes[child_id].edge_label = old_suffix.clone(); | |
| nodes[child_id].parent = Some(split_id); | |
| nodes[split_id].subtree_patterns = nodes[child_id].subtree_patterns; | |
| nodes[split_id].subtree_weight = nodes[child_id].subtree_weight; | |
| nodes[split_id] | |
| .children_by_label | |
| .insert(old_suffix, child_id); | |
| path.push(split_id); | |
| if new_suffix.is_empty() { | |
| return (split_id, path); | |
| } | |
| let new_child_id = push_node(nodes, split_id, new_suffix.clone()); | |
| nodes[split_id] | |
| .children_by_label | |
| .insert(new_suffix, new_child_id); | |
| path.push(new_child_id); | |
| return (new_child_id, path); | |
| } | |
| } | |
| fn push_node(nodes: &mut Vec<TrieNode>, parent_id: usize, edge_label: String) -> usize { | |
| let child_id = nodes.len(); | |
| nodes.push(TrieNode { | |
| id: child_id, | |
| edge_label, | |
| depth: nodes[parent_id].depth + 1, | |
| parent: Some(parent_id), | |
| children_by_label: BTreeMap::new(), | |
| subtree_patterns: 0, | |
| subtree_weight: 0, | |
| }); | |
| child_id | |
| } | |
| fn common_prefix_len(left: &str, right: &str) -> usize { | |
| let mut len = 0; | |
| for ((left_index, left_ch), (_, right_ch)) in left.char_indices().zip(right.char_indices()) { | |
| if left_ch != right_ch { | |
| break; | |
| } | |
| len = left_index + left_ch.len_utf8(); | |
| } | |
| len | |
| } | |
| fn assign_depths(nodes: &mut [TrieNode]) { | |
| let mut stack = vec![(0, 0)]; | |
| while let Some((node_id, depth)) = stack.pop() { | |
| nodes[node_id].depth = depth; | |
| let child_ids = nodes[node_id] | |
| .children_by_label | |
| .values() | |
| .copied() | |
| .collect::<Vec<_>>(); | |
| for child_id in child_ids { | |
| stack.push((child_id, depth + 1)); | |
| } | |
| } | |
| } | |
| fn find_episode_prefix(value: &str) -> Option<(String, String)> { | |
| let best_end = EPISODE_PATTERNS | |
| .iter() | |
| .filter_map(|pattern| pattern.find(value).map(|matched| matched.end())) | |
| .chain(find_delimited_number_episode_end(value)) | |
| .max()?; | |
| let prefix = value[..best_end].trim_end().to_owned(); | |
| let suffix = value[best_end..].trim().to_owned(); | |
| Some((prefix, suffix)) | |
| } | |
| fn find_delimited_number_episode_end(value: &str) -> Option<usize> { | |
| let mut digits_start = None; | |
| let mut digit_count = 0; | |
| for (index, ch) in value | |
| .char_indices() | |
| .chain(std::iter::once((value.len(), '\0'))) | |
| { | |
| if ch.is_ascii_digit() { | |
| if digits_start.is_none() { | |
| digits_start = Some(index); | |
| } | |
| digit_count += 1; | |
| continue; | |
| } | |
| if let Some(start) = digits_start { | |
| if (1..=4).contains(&digit_count) | |
| && has_episode_left_boundary(value, start) | |
| && has_episode_right_boundary(ch) | |
| { | |
| return Some(index); | |
| } | |
| } | |
| digits_start = None; | |
| digit_count = 0; | |
| } | |
| None | |
| } | |
| fn has_episode_left_boundary(value: &str, digits_start: usize) -> bool { | |
| if digits_start == 0 { | |
| return true; | |
| } | |
| value[..digits_start] | |
| .chars() | |
| .next_back() | |
| .is_some_and(|ch| ch.is_whitespace() || matches!(ch, '.' | '_' | '-')) | |
| } | |
| fn has_episode_right_boundary(ch: char) -> bool { | |
| ch == '\0' | |
| || ch.is_whitespace() | |
| || matches!(ch, '.' | '_' | '-' | ']' | ')' | '【' | '】' | '[') | |
| } | |
| fn digit_skeleton(text: &str) -> String { | |
| DIGITS.replace_all(text, "<NUM>").into_owned() | |
| } | |
| fn display_path(path: &Path) -> String { | |
| path.display().to_string() | |
| } | |
| fn default_count() -> usize { | |
| 1 | |
| } | |
| fn default_annotations_value() -> serde_json::Value { | |
| serde_json::json!({}) | |
| } | |
| mod tests { | |
| use super::*; | |
| fn dag_merges_equal_suffix_nodes_without_normalizing_digits() { | |
| let source = SourceGraph { | |
| meta: serde_json::json!({}), | |
| nodes: vec![ | |
| SourceNode { | |
| id: 0, | |
| edge_label: String::new(), | |
| depth: 0, | |
| children: vec![1, 3], | |
| }, | |
| SourceNode { | |
| id: 1, | |
| edge_label: "A01".to_owned(), | |
| depth: 1, | |
| children: vec![2], | |
| }, | |
| SourceNode { | |
| id: 2, | |
| edge_label: "X".to_owned(), | |
| depth: 2, | |
| children: vec![], | |
| }, | |
| SourceNode { | |
| id: 3, | |
| edge_label: "B02".to_owned(), | |
| depth: 1, | |
| children: vec![4], | |
| }, | |
| SourceNode { | |
| id: 4, | |
| edge_label: "X".to_owned(), | |
| depth: 2, | |
| children: vec![], | |
| }, | |
| ], | |
| terminals: vec![ | |
| SourceTerminal { | |
| terminal_id: None, | |
| node_id: 2, | |
| prefix: "A01X".to_owned(), | |
| digit_skeleton: "A<NUM>X".to_owned(), | |
| count: 1, | |
| weight: 1, | |
| uses_path_count: 0, | |
| has_trailing_hash_count: 0, | |
| suffix_examples: vec![], | |
| value_examples: vec!["A01X".to_owned()], | |
| annotations: serde_json::json!({}), | |
| }, | |
| SourceTerminal { | |
| terminal_id: None, | |
| node_id: 4, | |
| prefix: "B02X".to_owned(), | |
| digit_skeleton: "B<NUM>X".to_owned(), | |
| count: 1, | |
| weight: 1, | |
| uses_path_count: 0, | |
| has_trailing_hash_count: 0, | |
| suffix_examples: vec![], | |
| value_examples: vec!["B02X".to_owned()], | |
| annotations: serde_json::json!({}), | |
| }, | |
| ], | |
| }; | |
| let dag = minimize_source_graph( | |
| &source, | |
| Path::new("dmhy_prefix_graph.json"), | |
| Path::new("dmhy_prefix_dag.json"), | |
| ) | |
| .unwrap(); | |
| assert!(dag.meta.preserves_digits); | |
| assert!(dag.meta.merged_nodes >= 2); | |
| assert_eq!(dag.terminals[0].prefix, "A01X"); | |
| assert_eq!(dag.terminals[1].prefix, "B02X"); | |
| assert_eq!(dag.terminals[0].node_id, dag.terminals[1].node_id); | |
| let shared_parent_targets = dag.nodes[0] | |
| .children | |
| .iter() | |
| .filter(|edge| edge.label == "A01" || edge.label == "B02") | |
| .map(|edge| edge.target) | |
| .collect::<BTreeSet<_>>(); | |
| assert_eq!(shared_parent_targets.len(), 1); | |
| let shared_parent = shared_parent_targets.into_iter().next().unwrap(); | |
| assert_eq!(dag.nodes[shared_parent].children.len(), 1); | |
| assert_eq!(dag.nodes[shared_parent].children[0].label, "X"); | |
| assert_eq!(dag.nodes[shared_parent].reachable_terminals, 2); | |
| assert_eq!(dag.nodes[shared_parent].reachable_weight, 2); | |
| } | |
| } | |