morpheme 0.4.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

//! Replace a pattern in every token.

use serde::{Deserialize, Deserializer, Serialize};

use crate::error::Result;
use crate::pattern::{Pattern, SplitPattern, SysRegex};
use crate::traits::Decoder;

/// Replaces every match of `pattern` in each token with `content`
/// (e.g. `▁` → `" "` for SentencePiece models).
///
/// # Example
///
/// ```
/// use morpheme::decoders::Replace;
/// use morpheme::Decoder;
///
/// let d = Replace::new("▁", " ")?;
/// assert_eq!(d.decode(vec!["▁Hello".to_string(), "▁world".to_string()])?, " Hello world");
/// # Ok::<(), morpheme::Error>(())
/// ```
#[derive(Debug, Clone, Serialize)]
#[non_exhaustive]
pub struct Replace {
    /// What to replace.
    pub(crate) pattern: SplitPattern,
    /// The replacement.
    pub(crate) content: String,
    #[serde(skip)]
    regex: SysRegex,
}

impl PartialEq for Replace {
    fn eq(&self, other: &Self) -> bool {
        self.pattern == other.pattern && self.content == other.content
    }
}

impl<'de> Deserialize<'de> for Replace {
    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> std::result::Result<Self, D::Error> {
        #[derive(Deserialize)]
        struct Helper {
            pattern: SplitPattern,
            content: String,
        }
        let h = Helper::deserialize(deserializer)?;
        Replace::new(h.pattern, h.content).map_err(serde::de::Error::custom)
    }
}

impl Replace {
    /// Build a `Replace` decoder; fails if the regex does not compile.
    pub fn new(pattern: impl Into<SplitPattern>, content: impl Into<String>) -> Result<Self> {
        let pattern = pattern.into();
        let regex = pattern.to_regex()?;
        Ok(Self {
            pattern,
            content: content.into(),
            regex,
        })
    }

    /// The pattern being replaced.
    pub fn pattern(&self) -> &SplitPattern {
        &self.pattern
    }

    /// The replacement text.
    pub fn content(&self) -> &str {
        &self.content
    }
}

impl Decoder for Replace {
    fn decode_chain(&self, tokens: Vec<String>) -> Result<Vec<String>> {
        tokens
            .into_iter()
            .map(|token| {
                let mut out = String::with_capacity(token.len());
                for ((s, e), is_match) in (&self.regex).find_matches(&token)? {
                    if is_match {
                        out.push_str(&self.content);
                    } else {
                        out.push_str(&token[s..e]);
                    }
                }
                Ok(out)
            })
            .collect()
    }
}

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

    #[test]
    fn replaces() {
        let d = Replace::new(SplitPattern::String("▁".into()), " ").unwrap();
        let out = d
            .decode_chain(vec!["▁Hey".into(), "▁friend".into(), "!".into()])
            .unwrap();
        assert_eq!(out, vec![" Hey", " friend", "!"]);
    }
}