mtc-token-healing 0.3.3

Token healing implementation
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92

use std::collections::BTreeSet;

use crate::VocabPrefixAutomaton;

fn testcase_parse_chars<T: AsRef<str>>(
    automaton: &VocabPrefixAutomaton,
    vocab_sorted: &[&str],
    text: T,
) {
    let text = text.as_ref();
    let res: BTreeSet<_> = automaton
        .parse_bytes(text, 0)
        .into_iter()
        .map(|(pos, rank_range)| (pos, rank_range.lower, rank_range.upper))
        .collect();
    println!("{text}: {res:?}");
    for (pos, _) in text.char_indices() {
        let prefix = &text[pos..];
        let range = vocab_sorted
            .iter()
            .enumerate()
            .fold(None, |state, (k, token)| {
                if token.starts_with(prefix) {
                    if let Some((u, v)) = state {
                        assert_eq!(v, k);
                        Some((u, k + 1))
                    } else {
                        Some((k, k + 1))
                    }
                } else {
                    state
                }
            });
        if let Some((lower, upper)) = range {
            println!("{pos}: {:?}", &vocab_sorted[lower..upper]);
            assert!(res.contains(&(pos, lower as _, upper as _)));
        }
    }
}

fn testcase_vocab_prefix(vocab: &[&str], texts: &[&str]) {
    let automaton = VocabPrefixAutomaton::new(vocab);
    let vocab_sorted = automaton
        .order()
        .iter()
        .map(|&i| vocab[i as usize])
        .collect::<Vec<_>>();

    println!("vocab_sorted: {vocab_sorted:?}");

    for text in texts {
        testcase_parse_chars(&automaton, vocab_sorted.as_slice(), text);
    }
}

#[test]
fn test_chinese_vocab_prefix() {
    let vocab = ["歌曲", "聆听歌曲", "播放歌曲", "歌词", "查看歌词"];
    let texts = [
        "歌曲",
        "聆听歌曲",
        "聆听歌曲",
        "聆听歌曲",
        "播放歌曲",
        "播放歌曲",
        "播放歌曲",
        "歌词",
        "查看歌词",
        "查看歌词",
        "听歌曲",
        "听歌曲",
        "放歌曲",
        "听歌",
        "放歌",
        "词",
        "查看",
        "bba",
        "bbb",
        "bba",
        "bba",
        "cacab",
        "ccc",
    ];
    testcase_vocab_prefix(&vocab, &texts);
}

#[test]
fn test_simple_vocab_prefix() {
    let vocab = ["bb", "ca", "ab", "c", "aa", "bbaa", "a", "cc", "b"];
    let texts = ["bba", "bbb", "bba", "bba", "cacab", "ccc"];
    testcase_vocab_prefix(&vocab, &texts);
}