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

//! Whitespace pre-tokenizers.

use std::sync::OnceLock;

use crate::error::Result;
use crate::normalized_string::SplitDelimiterBehavior;
use crate::pattern::{Invert, SysRegex};
use crate::pre_tokenized_string::PreTokenizedString;
use crate::traits::PreTokenizer;

/// Splits into runs of word chars and runs of other non-space chars:
/// the regex `\w+|[^\w\s]+`. `"Hey man!"` → `["Hey", "man", "!"]`.
///
/// # 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::Whitespace;
///
/// // Words and runs of punctuation (`\w+|[^\w\s]+`).
/// assert_eq!(split(&Whitespace, "Hello, world!!"), ["Hello", ",", "world", "!!"]);
/// ```
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct Whitespace;

super::impl_unit_serde!(Whitespace);

fn word_regex() -> &'static SysRegex {
    static RE: OnceLock<SysRegex> = OnceLock::new();
    RE.get_or_init(|| SysRegex::new(r"\w+|[^\w\s]+").expect("valid regex"))
}

impl PreTokenizer for Whitespace {
    fn pre_tokenize(&self, pretokenized: &mut PreTokenizedString) -> Result<()> {
        pretokenized.split(|_, s| s.split(Invert(word_regex()), SplitDelimiterBehavior::Removed))
    }
}

/// Splits on whitespace only. `"Hey man!"` → `["Hey", "man!"]`.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct WhitespaceSplit;

super::impl_unit_serde!(WhitespaceSplit);

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

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

    #[test]
    fn whitespace() {
        assert_eq!(
            splits(&Whitespace, "Hey man!"),
            vec![("Hey", (0, 3)), ("man", (4, 7)), ("!", (7, 8))]
        );
        assert_eq!(
            splits(&Whitespace, "How are you doing?"),
            vec![
                ("How", (0, 3)),
                ("are", (4, 7)),
                ("you", (8, 11)),
                ("doing", (12, 17)),
                ("?", (17, 18)),
            ]
        );
        assert!(splits(&Whitespace, "\n").is_empty());
    }

    #[test]
    fn whitespace_split() {
        assert_eq!(
            splits(&WhitespaceSplit, "Hey, man, Good?"),
            vec![("Hey,", (0, 4)), ("man,", (5, 9)), ("Good?", (10, 15))]
        );
    }
}