diff --git a/CHANGELOG.md b/CHANGELOG.md index e65ecf3..35e852c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -41,6 +41,10 @@ Major release with breaking changes. ### Fixed - Fixed major slowdown risks on highly ambiguous tokenization paths. +- Fixed superlinear (quadratic) tokenization time in `newmm` on long text: + the ambiguity graph is no longer cleared with `HashMap::clear()` (which costs + O(capacity) per resolution), and BFS rebuilds its path from a predecessor map + instead of copying the path at every step. ### Migration diff --git a/src/tokenizer/newmm.rs b/src/tokenizer/newmm.rs index 007377a..1bd5c9d 100644 --- a/src/tokenizer/newmm.rs +++ b/src/tokenizer/newmm.rs @@ -306,32 +306,35 @@ impl NewmmTokenizer { graph: &HashMap>, start: CharacterIndex, goal: CharacterIndex, - current_queue: &mut VecDeque<(usize, Vec)>, + current_queue: &mut VecDeque, ) -> AnyResult> { current_queue.clear(); - // The visited set ensures each vertex is enqueued at most once. - // This avoids the exponential blow-up from revisiting the same nodes. - let mut visited: HashSet = - HashSet::with_capacity_and_hasher(goal - start, Default::default()); - visited.insert(start); + // Each vertex is enqueued at most once (this also avoids the + // exponential blow-up from revisiting nodes). Instead of copying the + // whole path at every step, record one predecessor per vertex and + // rebuild the path once, when the goal is reached. + let mut predecessor: HashMap = HashMap::default(); + predecessor.insert(start, start); + current_queue.push_back(start); - let mut init_path: Vec = Vec::with_capacity(goal - start); - init_path.push(start); - current_queue.push_back((start, init_path)); - - while let Some((vertex, path)) = current_queue.pop_front() { + while let Some(vertex) = current_queue.pop_front() { if let Some(idk) = graph.get(&vertex) { for &position in idk { if position == goal { - let mut appended_path = path; - appended_path.push(position); - return Ok(appended_path); - } else if !visited.contains(&position) { - visited.insert(position); - let mut appended_path = path.clone(); - appended_path.push(position); - current_queue.push_back((position, appended_path)); + let mut path = vec![goal, vertex]; + let mut current = vertex; + while current != start { + current = predecessor[¤t]; + path.push(current); + } + path.reverse(); + return Ok(path); + } else if let std::collections::hash_map::Entry::Vacant(entry) = + predecessor.entry(position) + { + entry.insert(vertex); + current_queue.push_back(position); } } }; @@ -343,10 +346,9 @@ impl NewmmTokenizer { fn one_cut<'a>(input: &'a CharString, custom_dict: &D) -> AnyResult> { let text = input; let input_char_len = text.chars_len(); - let mut reused_queue: VecDeque<(usize, Vec)> = VecDeque::with_capacity(10); + let mut reused_queue: VecDeque = VecDeque::with_capacity(10); let mut graph_size: usize = 0; let mut graph: HashMap> = HashMap::default(); - graph.reserve(input_char_len / 10); let mut result_str: Vec<&str> = Vec::with_capacity(input_char_len / 10); // All positions are character indices (not byte offsets). @@ -395,7 +397,9 @@ impl NewmmTokenizer { &mut reused_queue, )?; graph_size = 0; - graph.clear(); + // Not `clear()`: it costs O(capacity), which grows with the + // largest graph seen, and this runs once per resolution. + graph = HashMap::default(); for &position in group_of_end_position_candidate.iter().skip(1) { let token = text.substring_as_str(end_position, position); @@ -465,7 +469,7 @@ impl NewmmTokenizer { } graph_size = 0; - graph.clear(); + graph = HashMap::default(); let token = text.substring_as_str(begin_position, end_position); result_str.push(token); position_list.push(end_position); diff --git a/tests/test_tokenizer.rs b/tests/test_tokenizer.rs index 2a386d1..7906f4c 100644 --- a/tests/test_tokenizer.rs +++ b/tests/test_tokenizer.rs @@ -391,3 +391,53 @@ fn test_newmm_ambiguous_performance() { "tokenization took {elapsed:?}; BFS path explosion may still be present" ); } + +/// Best of three wall-clock timings of `segment` on `text`. +fn best_segment_time(tokenizer: &NewmmTokenizer, text: &str) -> std::time::Duration { + (0..3) + .map(|_| { + let start = std::time::Instant::now(); + let result = tokenizer.segment_with_options(text, false, None).unwrap(); + let elapsed = start.elapsed(); + assert_eq!(result.concat(), text); + elapsed + }) + .min() + .unwrap() +} + +/// Time must grow roughly linearly with text length, so a 4x longer text may +/// take about 4x as long; the 10x limit catches quadratic behavior (16x). +#[test] +fn test_newmm_persistent_ambiguity_scaling() { + let mut path = env!("CARGO_MANIFEST_DIR").to_string(); + path.push_str(DEFAULT_DICT_PATH); + let tokenizer = NewmmTokenizer::new(&path).unwrap(); + + // "กรรมกร" contains several overlapping dictionary words, so the + // ambiguity graph keeps growing across the whole text. + let small = best_segment_time(&tokenizer, &"กรรมกร".repeat(3_000)); + let large = best_segment_time(&tokenizer, &"กรรมกร".repeat(12_000)); + let limit = small.max(std::time::Duration::from_millis(10)) * 10; + assert!( + large < limit, + "4x longer text took {large:?} vs {small:?}; expected less than {limit:?}" + ); +} + +/// Ordinary long text must also scale linearly (10x longer text should take +/// about 10x as long; the 25x limit catches quadratic behavior). +#[test] +fn test_newmm_long_text_scaling() { + let mut path = env!("CARGO_MANIFEST_DIR").to_string(); + path.push_str(DEFAULT_DICT_PATH); + let tokenizer = NewmmTokenizer::new(&path).unwrap(); + + let small = best_segment_time(&tokenizer, &SECOND_TEXT.repeat(1_000)); + let large = best_segment_time(&tokenizer, &SECOND_TEXT.repeat(10_000)); + let limit = small.max(std::time::Duration::from_millis(10)) * 25; + assert!( + large < limit, + "10x longer text took {large:?} vs {small:?}; expected less than {limit:?}" + ); +}