morpheme 0.2.0

Fast, pure-Rust subword tokenization (BPE, WordPiece, Unigram) compatible with Hugging Face tokenizers.
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
93
94
95
96
97
98
99
100
101
102
103
104
105

//! BERT pre-tokenizer: split on whitespace, isolate punctuation.

use unicode_categories::UnicodeCategories;

use crate::error::Result;
use crate::normalized_string::SplitDelimiterBehavior;
use crate::pre_tokenized_string::PreTokenizedString;
use crate::traits::PreTokenizer;

/// BERT's definition of punctuation: ASCII punctuation (which includes
/// symbols like `$`, `^`, `` ` ``) plus every Unicode `P*` category.
pub(crate) fn is_bert_punc(c: char) -> bool {
    c.is_ascii_punctuation() || c.is_punctuation()
}

/// Splits on whitespace (removed) and isolates every punctuation char.
///
/// `"Hey friend!  How?!"` → `["Hey", "friend", "!", "How", "?", "!"]`.
///
/// # Example
///
/// ```
/// use morpheme::{OffsetType, PreTokenizedString, PreTokenizer};
///
/// fn split(pt: &impl PreTokenizer, text: &str) -> Vec<String> {
///     let mut s = PreTokenizedString::from(text);
///     pt.pre_tokenize(&mut s).unwrap();
///     s.get_splits(OffsetType::Byte).into_iter().map(|(w, _, _)| w.to_owned()).collect()
/// }
///
/// use morpheme::pre_tokenizers::BertPreTokenizer;
///
/// // Splits on whitespace and isolates every punctuation char.
/// assert_eq!(split(&BertPreTokenizer, "Hello, world!!"), ["Hello", ",", "world", "!", "!"]);
/// ```
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct BertPreTokenizer;

super::impl_unit_serde!(BertPreTokenizer);

impl PreTokenizer for BertPreTokenizer {
    fn pre_tokenize(&self, pretokenized: &mut PreTokenizedString) -> Result<()> {
        pretokenized.split(|_, s| s.split(char::is_whitespace, SplitDelimiterBehavior::Removed))?;
        pretokenized.split(|_, s| s.split(is_bert_punc, SplitDelimiterBehavior::Isolated))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::NormalizedString;
    use crate::pre_tokenizers::tests::splits;

    #[test]
    fn basic() {
        assert_eq!(
            splits(&BertPreTokenizer, "Hey friend!     How are you?!?"),
            vec![
                ("Hey", (0, 3)),
                ("friend", (4, 10)),
                ("!", (10, 11)),
                ("How", (16, 19)),
                ("are", (20, 23)),
                ("you", (24, 27)),
                ("?", (27, 28)),
                ("!", (28, 29)),
                ("?", (29, 30)),
            ]
        );
    }

    #[test]
    fn chinese_chars_padded_by_normalizer() {
        // Simulate BertNormalizer's handle_chinese_chars.
        let mut n = NormalizedString::from("野口里佳 Noguchi Rika");
        let dest: Vec<(char, isize)> = n
            .get()
            .chars()
            .flat_map(|c| {
                if (c as usize) > 0x4E00 {
                    vec![(' ', 0), (c, 1), (' ', 1)]
                } else {
                    vec![(c, 0)]
                }
            })
            .collect();
        n.transform(dest, 0);
        let mut pts = PreTokenizedString::from(n);
        BertPreTokenizer.pre_tokenize(&mut pts).unwrap();
        let got: Vec<_> = pts
            .get_splits(crate::OffsetType::Byte)
            .into_iter()
            .map(|(s, o, _)| (s.to_owned(), o))
            .collect();
        let want: Vec<(String, (usize, usize))> = vec![
            ("野".into(), (0, 3)),
            ("口".into(), (3, 6)),
            ("里".into(), (6, 9)),
            ("佳".into(), (9, 12)),
            ("Noguchi".into(), (13, 20)),
            ("Rika".into(), (21, 25)),
        ];
        assert_eq!(got, want);
    }
}