morpheme/models/
mod.rs

1//! Models: BPE, WordPiece, WordLevel and Unigram.
2
3pub(crate) mod bpe;
4pub(crate) mod unigram;
5pub(crate) mod wordlevel;
6pub(crate) mod wordpiece;
7
8use std::collections::HashMap;
9
10use serde::{Deserialize, Deserializer, Serialize};
11
12use crate::Token;
13use crate::error::Result;
14use crate::traits::Model;
15
16pub use bpe::{Bpe, BpeBuilder, Merges, Vocab};
17pub use unigram::Unigram;
18pub use wordlevel::{WordLevel, WordLevelBuilder};
19pub use wordpiece::{WordPiece, WordPieceBuilder};
20
21/// Any built-in model. This is what a [`crate::Tokenizer`] holds and
22/// what `tokenizer.json`'s `"model"` field (de)serializes to.
23#[derive(Debug, Clone, PartialEq, Serialize)]
24#[serde(untagged)]
25#[non_exhaustive]
26pub enum ModelWrapper {
27    /// Byte-Pair Encoding.
28    Bpe(Bpe),
29    /// WordPiece (BERT).
30    WordPiece(WordPiece),
31    /// Whole-word lookup.
32    WordLevel(WordLevel),
33    /// Unigram language model (SentencePiece).
34    Unigram(Unigram),
35}
36
37impl<'de> Deserialize<'de> for ModelWrapper {
38    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> std::result::Result<Self, D::Error> {
39        use serde::de::Error;
40        let value = serde_json::Value::deserialize(deserializer)?;
41        let ty = value
42            .get("type")
43            .and_then(|t| t.as_str())
44            .map(str::to_owned)
45            .or_else(|| infer_model_type(&value).map(str::to_owned))
46            .ok_or_else(|| D::Error::custom("cannot determine the model type"))?;
47        let model = match ty.as_str() {
48            "BPE" => ModelWrapper::Bpe(serde_json::from_value(value).map_err(D::Error::custom)?),
49            "WordPiece" => {
50                ModelWrapper::WordPiece(serde_json::from_value(value).map_err(D::Error::custom)?)
51            }
52            "WordLevel" => {
53                ModelWrapper::WordLevel(serde_json::from_value(value).map_err(D::Error::custom)?)
54            }
55            "Unigram" => {
56                ModelWrapper::Unigram(serde_json::from_value(value).map_err(D::Error::custom)?)
57            }
58            other => {
59                return Err(D::Error::custom(format!(
60                    "unsupported model type {other:?} (expected BPE, WordPiece, WordLevel or Unigram)"
61                )));
62            }
63        };
64        Ok(model)
65    }
66}
67
68/// Files written by older versions of `tokenizers` omit the model's
69/// `"type"`. Infer it from the fields present.
70fn infer_model_type(value: &serde_json::Value) -> Option<&'static str> {
71    let obj = value.as_object()?;
72    if obj.contains_key("merges") {
73        Some("BPE")
74    } else if obj.contains_key("max_input_chars_per_word")
75        || obj.contains_key("continuing_subword_prefix")
76    {
77        Some("WordPiece")
78    } else if obj.contains_key("unk_id") || obj.get("vocab").is_some_and(|v| v.is_array()) {
79        Some("Unigram")
80    } else if obj.contains_key("vocab") {
81        Some("WordLevel")
82    } else {
83        None
84    }
85}
86
87macro_rules! dispatch {
88    ($self:ident, $m:ident => $e:expr) => {
89        match $self {
90            ModelWrapper::Bpe($m) => $e,
91            ModelWrapper::WordPiece($m) => $e,
92            ModelWrapper::WordLevel($m) => $e,
93            ModelWrapper::Unigram($m) => $e,
94        }
95    };
96}
97
98impl Model for ModelWrapper {
99    fn tokenize(&self, sequence: &str) -> Result<Vec<Token>> {
100        dispatch!(self, m => m.tokenize(sequence))
101    }
102    fn token_to_id(&self, token: &str) -> Option<u32> {
103        dispatch!(self, m => m.token_to_id(token))
104    }
105    fn id_to_token(&self, id: u32) -> Option<String> {
106        dispatch!(self, m => m.id_to_token(id))
107    }
108    fn vocab(&self) -> HashMap<String, u32> {
109        dispatch!(self, m => m.vocab())
110    }
111    fn vocab_size(&self) -> usize {
112        dispatch!(self, m => m.vocab_size())
113    }
114}
115
116impl From<Bpe> for ModelWrapper {
117    fn from(m: Bpe) -> Self {
118        ModelWrapper::Bpe(m)
119    }
120}
121
122impl From<WordPiece> for ModelWrapper {
123    fn from(m: WordPiece) -> Self {
124        ModelWrapper::WordPiece(m)
125    }
126}
127
128impl From<WordLevel> for ModelWrapper {
129    fn from(m: WordLevel) -> Self {
130        ModelWrapper::WordLevel(m)
131    }
132}
133
134impl From<Unigram> for ModelWrapper {
135    fn from(m: Unigram) -> Self {
136        ModelWrapper::Unigram(m)
137    }
138}
morpheme/models/mod.rs

morpheme/models/
mod.rs
