ModerRAS's picture
Add DMHY prefix graph annotation workflow
33bb11c
Raw
History Blame Contribute Delete
32.2 kB
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());
#[derive(Parser)]
#[command(about = "Build a deterministic DMHY episode-prefix trie graph")]
struct Args {
#[arg(long, default_value = "datasets\\AnimeName\\dmhy_list.jsonl")]
input: PathBuf,
#[arg(long, default_value = "datasets\\AnimeName\\dmhy_prefix_graph.json")]
output: PathBuf,
#[arg(long, default_value_t = 2)]
min_count: usize,
#[arg(long, default_value_t = 5)]
example_count: usize,
#[arg(long)]
from_prefix_graph: Option<PathBuf>,
#[arg(long)]
dag_output: Option<PathBuf>,
}
#[derive(Clone, Deserialize)]
struct InputRecord {
value: String,
#[serde(default)]
uses_path: bool,
#[serde(default)]
has_trailing_hash: bool,
#[serde(default, rename = "has_digits")]
_has_digits: bool,
#[serde(default)]
#[serde(rename = "digit_skeleton")]
_digit_skeleton: String,
#[serde(default = "default_count")]
count: usize,
}
#[derive(Clone)]
struct PatternObservation {
source_index: usize,
prefix: String,
digit_skeleton: String,
suffix: String,
value: String,
uses_path: bool,
has_trailing_hash: bool,
count: usize,
}
#[derive(Default)]
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,
}
#[derive(Serialize)]
struct GraphOutput {
meta: Meta,
nodes: Vec<OutputNode>,
terminals: Vec<OutputTerminal>,
}
#[derive(Serialize)]
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,
}
#[derive(Serialize)]
struct TokenizerMeta {
version: &'static str,
notes: Vec<&'static str>,
}
#[derive(Serialize)]
struct OutputNode {
id: usize,
edge_label: String,
depth: usize,
children: Vec<usize>,
subtree_patterns: usize,
subtree_weight: usize,
}
#[derive(Serialize)]
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,
}
#[derive(Default, Serialize)]
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>,
}
#[derive(Deserialize)]
struct SourceGraph {
#[allow(dead_code)]
meta: serde_json::Value,
nodes: Vec<SourceNode>,
terminals: Vec<SourceTerminal>,
}
#[derive(Deserialize)]
struct SourceNode {
id: usize,
edge_label: String,
#[allow(dead_code)]
depth: usize,
#[serde(default)]
children: Vec<usize>,
}
#[derive(Clone, Deserialize, Serialize)]
struct SourceTerminal {
#[serde(default, skip_serializing_if = "Option::is_none")]
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>,
#[serde(default = "default_annotations_value")]
annotations: serde_json::Value,
}
#[derive(Serialize)]
struct DagOutput {
meta: DagMeta,
root: usize,
nodes: Vec<DagNode>,
terminals: Vec<DagTerminal>,
}
#[derive(Serialize)]
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>,
}
#[derive(Clone, Serialize)]
struct DagNode {
id: usize,
terminal: bool,
children: Vec<DagEdge>,
incoming_count: usize,
reachable_terminals: usize,
reachable_weight: usize,
}
#[derive(Clone, Serialize)]
struct DagEdge {
label: String,
target: usize,
}
#[derive(Serialize)]
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,
}
#[derive(Clone, Eq, Ord, PartialEq, PartialOrd)]
struct NodeSignature {
terminal: bool,
children: Vec<(String, usize)>,
}
#[derive(Clone)]
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!({})
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
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);
}
}