1use rayon::prelude::*;
9use regex::Regex;
10use std::collections::HashMap;
11use std::sync::LazyLock;
12
13pub type PositionedDecodeResult = (
15 Vec<String>,
16 Vec<f32>,
17 Vec<Vec<f32>>,
18 Vec<Vec<usize>>,
19 Vec<usize>,
20);
21
22#[derive(Debug, PartialEq)]
28pub(crate) struct CTCArgmaxOutput {
29 batch_size: usize,
30 sequence_length: usize,
31 indices: Vec<usize>,
32 probabilities: Vec<f32>,
33}
34
35static ALPHANUMERIC_REGEX: LazyLock<Regex> = LazyLock::new(|| {
36 Regex::new(r"[a-zA-Z0-9 :*./%+-]").expect("static regex: alphanumeric decoder pattern")
37});
38
39#[inline]
46fn argmax_row(row: ndarray::ArrayView1<f32>) -> Option<(usize, f32)> {
47 match row.as_slice() {
48 Some(slice) => crate::processors::simd::argmax(slice),
49 None => row
50 .iter()
51 .copied()
52 .enumerate()
53 .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)),
54 }
55}
56
57pub struct BaseRecLabelDecode {
68 reverse: bool,
69 dict: HashMap<char, usize>,
70 character: Vec<char>,
71}
72
73impl BaseRecLabelDecode {
74 pub fn new(character_str: Option<&str>, use_space_char: bool) -> Self {
84 let mut character_list: Vec<char> = if let Some(chars) = character_str {
85 chars.chars().collect()
86 } else {
87 "0123456789abcdefghijklmnopqrstuvwxyz".chars().collect()
88 };
89
90 if use_space_char {
91 character_list.push(' ');
92 }
93
94 character_list = Self::add_special_char(character_list);
95
96 let mut dict = HashMap::new();
97 for (i, &char) in character_list.iter().enumerate() {
98 dict.insert(char, i);
99 }
100
101 Self {
102 reverse: false,
103 dict,
104 character: character_list,
105 }
106 }
107
108 pub fn from_string_list(character_list: Option<&[String]>, use_space_char: bool) -> Self {
119 let mut chars: Vec<char> = if let Some(list) = character_list {
120 list.iter().filter_map(|s| s.chars().next()).collect()
121 } else {
122 "0123456789abcdefghijklmnopqrstuvwxyz".chars().collect()
123 };
124
125 if use_space_char {
126 chars.push(' ');
127 }
128
129 chars = Self::add_special_char(chars);
130
131 let mut dict = HashMap::new();
132 for (i, &char) in chars.iter().enumerate() {
133 dict.insert(char, i);
134 }
135
136 Self {
137 reverse: false,
138 dict,
139 character: chars,
140 }
141 }
142
143 fn pred_reverse(&self, pred: &str) -> String {
151 let mut pred_re = Vec::new();
152 let mut c_current = String::new();
153
154 for c in pred.chars() {
155 if !ALPHANUMERIC_REGEX.is_match(&c.to_string()) {
156 if !c_current.is_empty() {
157 pred_re.push(c_current.clone());
158 c_current.clear();
159 }
160 pred_re.push(c.to_string());
161 } else {
162 c_current.push(c);
163 }
164 }
165
166 if !c_current.is_empty() {
167 pred_re.push(c_current);
168 }
169
170 pred_re.reverse();
171 pred_re.join("")
172 }
173
174 fn add_special_char(character_list: Vec<char>) -> Vec<char> {
185 character_list
186 }
187
188 fn get_ignored_tokens(&self) -> Vec<usize> {
193 vec![self.get_blank_idx()]
194 }
195
196 pub fn decode(
206 &self,
207 text_index: &[Vec<usize>],
208 text_prob: Option<&[Vec<f32>]>,
209 is_remove_duplicate: bool,
210 ) -> Vec<(String, f32)> {
211 let mut result_list = Vec::new();
212 let ignored_tokens = self.get_ignored_tokens();
213
214 for (batch_idx, indices) in text_index.iter().enumerate() {
215 let mut selection = vec![true; indices.len()];
216
217 if is_remove_duplicate && indices.len() > 1 {
218 for i in 1..indices.len() {
219 if indices[i] == indices[i - 1] {
220 selection[i] = false;
221 }
222 }
223 }
224
225 for &ignored_token in &ignored_tokens {
226 for (i, &idx) in indices.iter().enumerate() {
227 if idx == ignored_token {
228 selection[i] = false;
229 }
230 }
231 }
232
233 let char_list: Vec<char> = indices
234 .iter()
235 .enumerate()
236 .filter(|(i, _)| selection[*i])
237 .filter_map(|(_, &text_id)| self.character.get(text_id).copied())
238 .collect();
239
240 let conf_list: Vec<f32> = if let Some(probs) = text_prob {
241 if batch_idx < probs.len() {
242 probs[batch_idx]
243 .iter()
244 .enumerate()
245 .filter(|(i, _)| *i < selection.len() && selection[*i])
246 .map(|(_, &prob)| prob)
247 .collect()
248 } else {
249 vec![1.0; char_list.len()]
250 }
251 } else {
252 vec![1.0; char_list.len()]
253 };
254
255 let conf_list = if conf_list.is_empty() {
256 vec![0.0]
257 } else {
258 conf_list
259 };
260
261 let mut text: String = char_list.iter().collect();
262
263 if self.reverse {
264 text = self.pred_reverse(&text);
265 }
266
267 let mean_conf = conf_list.iter().sum::<f32>() / conf_list.len() as f32;
268 result_list.push((text, mean_conf));
269 }
270
271 result_list
272 }
273
274 pub fn apply(&self, pred: &ndarray::Array3<f32>) -> (Vec<String>, Vec<f32>) {
285 if pred.is_empty() {
286 return (Vec::new(), Vec::new());
287 }
288
289 let batch_size = pred.shape()[0];
290 let mut all_texts = Vec::new();
291 let mut all_scores = Vec::new();
292
293 for batch_idx in 0..batch_size {
294 let preds = pred.index_axis(ndarray::Axis(0), batch_idx);
295
296 let mut sequence_idx = Vec::new();
297 let mut sequence_prob = Vec::new();
298
299 for row in preds.outer_iter() {
300 if let Some((idx, prob)) = argmax_row(row) {
301 sequence_idx.push(idx);
302 sequence_prob.push(prob);
303 } else {
304 sequence_idx.push(0);
305 sequence_prob.push(0.0);
306 }
307 }
308
309 let text = self.decode(&[sequence_idx], Some(&[sequence_prob]), true);
310
311 for (t, score) in text {
312 all_texts.push(t);
313 all_scores.push(score);
314 }
315 }
316
317 (all_texts, all_scores)
318 }
319
320 fn get_blank_idx(&self) -> usize {
325 0
326 }
327}
328
329pub struct CTCLabelDecode {
338 base: BaseRecLabelDecode,
339 blank_index: usize,
340}
341
342impl std::fmt::Debug for CTCLabelDecode {
343 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
344 f.debug_struct("CTCLabelDecode")
345 .field("character_count", &self.base.character.len())
346 .field("reverse", &self.base.reverse)
347 .finish()
348 }
349}
350
351impl CTCLabelDecode {
352 pub fn new(character_list: Option<&str>, use_space_char: bool) -> Self {
362 let mut base = BaseRecLabelDecode::new(character_list, use_space_char);
363
364 let mut new_character = vec!['\0'];
366 new_character.extend(base.character);
367
368 let mut new_dict = HashMap::new();
369 for (i, &char) in new_character.iter().enumerate() {
370 new_dict.insert(char, i);
371 }
372
373 base.character = new_character;
374 base.dict = new_dict;
375
376 let blank_index = 0;
377
378 Self { base, blank_index }
379 }
380
381 pub fn from_string_list(
393 character_list: Option<&[String]>,
394 use_space_char: bool,
395 has_explicit_blank: bool,
396 ) -> Self {
397 if has_explicit_blank {
398 let base = BaseRecLabelDecode::from_string_list(character_list, use_space_char);
399 Self {
400 base,
401 blank_index: 0,
402 }
403 } else {
404 let mut base = BaseRecLabelDecode::from_string_list(character_list, use_space_char);
405
406 let mut new_character = vec!['\0'];
408 new_character.extend(base.character);
409
410 let mut new_dict = HashMap::new();
411 for (i, &char) in new_character.iter().enumerate() {
412 new_dict.insert(char, i);
413 }
414
415 base.character = new_character;
416 base.dict = new_dict;
417
418 Self {
419 base,
420 blank_index: 0,
421 }
422 }
423 }
424
425 pub fn get_blank_index(&self) -> usize {
430 self.blank_index
431 }
432
433 pub fn get_character_list(&self) -> &[char] {
438 &self.base.character
439 }
440
441 pub fn get_character_count(&self) -> usize {
446 self.base.character.len()
447 }
448
449 pub(crate) fn argmax_predictions<S>(
453 &self,
454 pred: &ndarray::ArrayBase<S, ndarray::Ix3>,
455 ) -> CTCArgmaxOutput
456 where
457 S: ndarray::Data<Elem = f32> + Sync,
458 {
459 let [batch_size, sequence_length, vocab_size] = pred
460 .shape()
461 .try_into()
462 .expect("CTC predictions always have three dimensions");
463 let row_count = batch_size * sequence_length;
464
465 if pred.is_empty() {
468 return CTCArgmaxOutput {
469 batch_size: 0,
470 sequence_length: 0,
471 indices: Vec::new(),
472 probabilities: Vec::new(),
473 };
474 }
475
476 let (indices, probabilities): (Vec<usize>, Vec<f32>) = if let Some(data) = pred.as_slice() {
480 data.par_chunks_exact(vocab_size)
481 .map(|row| crate::processors::simd::argmax(row).unwrap_or((self.blank_index, 0.0)))
482 .unzip()
483 } else {
484 (0..row_count)
485 .into_par_iter()
486 .map(|row_idx| {
487 let batch_idx = row_idx / sequence_length;
488 let time_idx = row_idx % sequence_length;
489 argmax_row(pred.slice(ndarray::s![batch_idx, time_idx, ..]))
490 .unwrap_or((self.blank_index, 0.0))
491 })
492 .unzip()
493 };
494
495 CTCArgmaxOutput {
496 batch_size,
497 sequence_length,
498 indices,
499 probabilities,
500 }
501 }
502
503 pub(crate) fn decode_argmax(&self, argmax: &CTCArgmaxOutput) -> (Vec<String>, Vec<f32>) {
506 let (all_texts, all_scores): (Vec<String>, Vec<f32>) = (0..argmax.batch_size)
507 .into_par_iter()
508 .map(|batch_idx| {
509 let start = batch_idx * argmax.sequence_length;
510 let end = start + argmax.sequence_length;
511 let sequence_idx = &argmax.indices[start..end];
512 let sequence_prob = &argmax.probabilities[start..end];
513
514 let mut filtered_prob = Vec::with_capacity(argmax.sequence_length);
515 let mut text = String::with_capacity(argmax.sequence_length);
516 let mut prev_idx = self.blank_index;
517 for (i, &idx) in sequence_idx.iter().enumerate() {
518 if idx != self.blank_index
519 && idx != prev_idx
520 && let Some(&ch) = self.base.character.get(idx)
521 {
522 text.push(ch);
523 filtered_prob.push(sequence_prob[i]);
524 }
525 prev_idx = idx;
526 }
527
528 let mean_conf = if filtered_prob.is_empty() {
529 0.0
530 } else {
531 filtered_prob.iter().sum::<f32>() / filtered_prob.len() as f32
532 };
533
534 (text, mean_conf)
535 })
536 .unzip();
537
538 (all_texts, all_scores)
539 }
540
541 pub(crate) fn decode_argmax_with_positions(
544 &self,
545 argmax: &CTCArgmaxOutput,
546 ) -> PositionedDecodeResult {
547 type PerItem = (String, f32, Vec<f32>, Vec<usize>, usize);
548 let per: Vec<PerItem> = (0..argmax.batch_size)
549 .into_par_iter()
550 .map(|batch_idx| {
551 let start = batch_idx * argmax.sequence_length;
552 let end = start + argmax.sequence_length;
553 let sequence_idx = &argmax.indices[start..end];
554 let sequence_prob = &argmax.probabilities[start..end];
555
556 let mut filtered_prob = Vec::with_capacity(argmax.sequence_length);
557 let mut filtered_timesteps = Vec::with_capacity(argmax.sequence_length);
558 let mut char_list = Vec::with_capacity(argmax.sequence_length);
559 let mut prev_idx = self.blank_index;
560 for (i, &idx) in sequence_idx.iter().enumerate() {
561 if idx != self.blank_index
562 && idx != prev_idx
563 && let Some(&ch) = self.base.character.get(idx)
564 {
565 char_list.push(ch);
566 filtered_prob.push(sequence_prob[i]);
567 filtered_timesteps.push(i);
568 }
569 prev_idx = idx;
570 }
571
572 let mean_conf = if filtered_prob.is_empty() {
573 0.0
574 } else {
575 filtered_prob.iter().sum::<f32>() / filtered_prob.len() as f32
576 };
577 let seq_len = argmax.sequence_length as f32;
578 let char_positions = filtered_timesteps
579 .iter()
580 .map(|×tep| timestep as f32 / seq_len)
581 .collect();
582 let text = char_list.iter().collect();
583
584 (
585 text,
586 mean_conf,
587 char_positions,
588 filtered_timesteps,
589 argmax.sequence_length,
590 )
591 })
592 .collect();
593
594 let mut all_texts = Vec::with_capacity(argmax.batch_size);
595 let mut all_scores = Vec::with_capacity(argmax.batch_size);
596 let mut all_positions = Vec::with_capacity(argmax.batch_size);
597 let mut all_col_indices = Vec::with_capacity(argmax.batch_size);
598 let mut all_seq_lengths = Vec::with_capacity(argmax.batch_size);
599 for (text, score, pos, cols, seq_len) in per {
600 all_texts.push(text);
601 all_scores.push(score);
602 all_positions.push(pos);
603 all_col_indices.push(cols);
604 all_seq_lengths.push(seq_len);
605 }
606
607 (
608 all_texts,
609 all_scores,
610 all_positions,
611 all_col_indices,
612 all_seq_lengths,
613 )
614 }
615
616 pub fn apply_with_positions<S>(
633 &self,
634 pred: &ndarray::ArrayBase<S, ndarray::Ix3>,
635 ) -> PositionedDecodeResult
636 where
637 S: ndarray::Data<Elem = f32> + Sync,
638 {
639 if pred.is_empty() {
640 return (Vec::new(), Vec::new(), Vec::new(), Vec::new(), Vec::new());
641 }
642 let argmax = self.argmax_predictions(pred);
643 self.decode_argmax_with_positions(&argmax)
644 }
645
646 pub fn apply<S>(&self, pred: &ndarray::ArrayBase<S, ndarray::Ix3>) -> (Vec<String>, Vec<f32>)
663 where
664 S: ndarray::Data<Elem = f32> + Sync,
665 {
666 if pred.is_empty() {
667 return (Vec::new(), Vec::new());
668 }
669 let argmax = self.argmax_predictions(pred);
670 self.decode_argmax(&argmax)
671 }
672}
673
674#[cfg(test)]
675mod tests {
676 use super::*;
677 use ndarray::Array3;
678
679 fn logits_with_winners(winners: &[&[(usize, f32)]], vocab_size: usize) -> Array3<f32> {
680 let batch_size = winners.len();
681 let sequence_length = winners.first().map_or(0, |sequence| sequence.len());
682 let mut logits = Array3::from_elem((batch_size, sequence_length, vocab_size), -10.0);
683 for (batch_idx, sequence) in winners.iter().enumerate() {
684 assert_eq!(sequence.len(), sequence_length);
685 for (time_idx, &(token_idx, probability)) in sequence.iter().enumerate() {
686 logits[[batch_idx, time_idx, token_idx]] = probability;
687 }
688 }
689 logits
690 }
691
692 #[test]
693 fn compact_argmax_preserves_ctc_text_scores_and_positions() {
694 let characters = vec!["a".to_string(), "b".to_string(), "c".to_string()];
695 let decoder = CTCLabelDecode::from_string_list(Some(&characters), false, false);
696 let logits = logits_with_winners(
698 &[
699 &[
700 (0, 0.9),
701 (1, 0.8),
702 (1, 0.7),
703 (0, 0.6),
704 (1, 0.5),
705 (2, 0.4),
706 (2, 0.3),
707 ],
708 &[
709 (3, 0.95),
710 (3, 0.85),
711 (4, 0.75),
712 (3, 0.65),
713 (0, 0.55),
714 (2, 0.45),
715 (0, 0.35),
716 ],
717 ],
718 5,
719 );
720
721 let argmax = decoder.argmax_predictions(&logits);
722 assert_eq!(argmax.batch_size, 2);
723 assert_eq!(argmax.sequence_length, 7);
724 assert_eq!(argmax.indices.len(), 14);
725 assert_eq!(argmax.probabilities.len(), 14);
726
727 let (texts, scores) = decoder.decode_argmax(&argmax);
728 assert_eq!(texts, ["aab", "ccb"]);
729 assert_eq!(
730 scores,
731 [(0.8 + 0.5 + 0.4) / 3.0, (0.95 + 0.65 + 0.45) / 3.0]
732 );
733
734 let (texts, scores, positions, columns, lengths) =
735 decoder.decode_argmax_with_positions(&argmax);
736 assert_eq!(texts, ["aab", "ccb"]);
737 assert_eq!(
738 scores,
739 [(0.8 + 0.5 + 0.4) / 3.0, (0.95 + 0.65 + 0.45) / 3.0]
740 );
741 assert_eq!(columns, [vec![1, 4, 5], vec![0, 3, 5]]);
742 assert_eq!(positions[0], [1.0 / 7.0, 4.0 / 7.0, 5.0 / 7.0]);
743 assert_eq!(positions[1], [0.0, 3.0 / 7.0, 5.0 / 7.0]);
744 assert_eq!(lengths, [7, 7]);
745 }
746
747 #[test]
748 fn compact_argmax_preserves_empty_tensor_behavior() {
749 let decoder = CTCLabelDecode::new(None, false);
750 let logits = Array3::<f32>::zeros((2, 0, decoder.get_character_count()));
751 let argmax = decoder.argmax_predictions(&logits);
752
753 assert_eq!(decoder.decode_argmax(&argmax), (Vec::new(), Vec::new()));
754 assert_eq!(
755 decoder.decode_argmax_with_positions(&argmax),
756 (Vec::new(), Vec::new(), Vec::new(), Vec::new(), Vec::new())
757 );
758 }
759}