sentencepiece_rust/
model.rs

1//! Parsed `ModelProto` — the subset of fields needed for inference.
2//!
3//! Field numbers and defaults mirror
4//! `Reference/sentencepiece/src/sentencepiece_model.proto`.
5
6use crate::error::{Error, Result};
7use crate::proto::Reader;
8
9/// Segmentation algorithm stored in the model (`TrainerSpec.model_type`).
10#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11pub enum ModelType {
12    Unigram,
13    Bpe,
14    Word,
15    Char,
16}
17
18/// Piece category (`ModelProto.SentencePiece.Type`).
19#[derive(Debug, Clone, Copy, PartialEq, Eq)]
20pub enum PieceType {
21    Normal,
22    Unknown,
23    Control,
24    UserDefined,
25    Byte,
26    Unused,
27}
28
29impl PieceType {
30    fn from_i32(v: i32) -> PieceType {
31        match v {
32            1 => PieceType::Normal,
33            2 => PieceType::Unknown,
34            3 => PieceType::Control,
35            4 => PieceType::UserDefined,
36            6 => PieceType::Byte,
37            5 => PieceType::Unused,
38            _ => PieceType::Normal,
39        }
40    }
41}
42
43#[derive(Debug, Clone)]
44pub struct SentencePiece {
45    pub piece: String,
46    pub score: f32,
47    pub kind: PieceType,
48}
49
50/// `NormalizerSpec` — text normalisation parameters.
51#[derive(Debug, Clone)]
52pub struct NormalizerSpec {
53    /// Compiled Darts double-array charsmap (empty for identity normalisation).
54    pub precompiled_charsmap: Vec<u8>,
55    pub add_dummy_prefix: bool,
56    pub remove_extra_whitespaces: bool,
57    pub escape_whitespaces: bool,
58}
59
60impl Default for NormalizerSpec {
61    fn default() -> Self {
62        NormalizerSpec {
63            precompiled_charsmap: Vec::new(),
64            add_dummy_prefix: true,
65            remove_extra_whitespaces: true,
66            escape_whitespaces: true,
67        }
68    }
69}
70
71/// The decoded model: vocabulary pieces plus the flags inference depends on.
72#[derive(Debug, Clone)]
73pub struct ModelProto {
74    pub pieces: Vec<SentencePiece>,
75    pub model_type: ModelType,
76    pub byte_fallback: bool,
77    pub treat_whitespace_as_suffix: bool,
78    pub unk_id: i32,
79    pub bos_id: i32,
80    pub eos_id: i32,
81    pub pad_id: i32,
82    /// Surface emitted for `<unk>` on decode (`TrainerSpec.unk_surface`).
83    pub unk_surface: String,
84    pub normalizer: NormalizerSpec,
85}
86
87impl ModelProto {
88    /// Parse a serialised `.model` file.
89    pub fn parse(bytes: &[u8]) -> Result<ModelProto> {
90        let mut model = ModelProto {
91            pieces: Vec::new(),
92            model_type: ModelType::Unigram,
93            byte_fallback: false,
94            treat_whitespace_as_suffix: false,
95            unk_id: 0,
96            bos_id: 1,
97            eos_id: 2,
98            pad_id: -1,
99            // proto2 default for unk_surface: " \xE2\x81\x87 " (space, U+2047, space).
100            unk_surface: " \u{2047} ".to_string(),
101            normalizer: NormalizerSpec::default(),
102        };
103
104        let mut r = Reader::new(bytes);
105        while let Some((number, value)) = r.next()? {
106            match number {
107                // repeated SentencePiece pieces = 1;
108                1 => {
109                    if let Some(b) = value.as_bytes() {
110                        model.pieces.push(parse_piece(b)?);
111                    }
112                }
113                // optional TrainerSpec trainer_spec = 2;
114                2 => {
115                    if let Some(b) = value.as_bytes() {
116                        parse_trainer_spec(b, &mut model)?;
117                    }
118                }
119                // optional NormalizerSpec normalizer_spec = 3;
120                3 => {
121                    if let Some(b) = value.as_bytes() {
122                        model.normalizer = parse_normalizer_spec(b)?;
123                    }
124                }
125                _ => {} // self_test_data, denormalizer_spec, extensions: ignored
126            }
127        }
128
129        if model.pieces.is_empty() {
130            return Err(Error::Model("model contains no pieces".into()));
131        }
132        Ok(model)
133    }
134}
135
136fn parse_piece(bytes: &[u8]) -> Result<SentencePiece> {
137    let mut piece = String::new();
138    let mut score = 0.0f32;
139    let mut kind = PieceType::Normal;
140    let mut r = Reader::new(bytes);
141    while let Some((number, value)) = r.next()? {
142        match number {
143            1 => piece = value.as_str().unwrap_or_default(),
144            2 => score = value.as_f32().unwrap_or(0.0),
145            3 => kind = PieceType::from_i32(value.as_i32().unwrap_or(1)),
146            _ => {}
147        }
148    }
149    Ok(SentencePiece { piece, score, kind })
150}
151
152fn parse_normalizer_spec(bytes: &[u8]) -> Result<NormalizerSpec> {
153    let mut spec = NormalizerSpec::default();
154    let mut r = Reader::new(bytes);
155    while let Some((number, value)) = r.next()? {
156        match number {
157            2 => spec.precompiled_charsmap = value.as_bytes().unwrap_or_default().to_vec(),
158            3 => spec.add_dummy_prefix = value.as_bool().unwrap_or(true),
159            4 => spec.remove_extra_whitespaces = value.as_bool().unwrap_or(true),
160            5 => spec.escape_whitespaces = value.as_bool().unwrap_or(true),
161            _ => {}
162        }
163    }
164    Ok(spec)
165}
166
167fn parse_trainer_spec(bytes: &[u8], model: &mut ModelProto) -> Result<()> {
168    let mut r = Reader::new(bytes);
169    while let Some((number, value)) = r.next()? {
170        match number {
171            3 => {
172                model.model_type = match value.as_i32().unwrap_or(1) {
173                    2 => ModelType::Bpe,
174                    3 => ModelType::Word,
175                    4 => ModelType::Char,
176                    _ => ModelType::Unigram,
177                }
178            }
179            24 => model.treat_whitespace_as_suffix = value.as_bool().unwrap_or(false),
180            35 => model.byte_fallback = value.as_bool().unwrap_or(false),
181            40 => model.unk_id = value.as_i32().unwrap_or(0),
182            41 => model.bos_id = value.as_i32().unwrap_or(1),
183            42 => model.eos_id = value.as_i32().unwrap_or(2),
184            43 => model.pad_id = value.as_i32().unwrap_or(-1),
185            44 => {
186                if let Some(s) = value.as_str() {
187                    model.unk_surface = s;
188                }
189            }
190            _ => {}
191        }
192    }
193    Ok(())
194}
sentencepiece_rust/model.rs

sentencepiece_rust/
model.rs
