laya/
tokenizer.rs

1//! Checkpoint tokenizer: a thin wrapper over the Hugging Face `tokenizers` crate that loads
2//! `tokenizer.json` directly and resolves Laya's special token ids. Mirrors the special-id
3//! resolution in `laya-ts/src/tokenizer.ts`.
4
5use crate::common::Tokenizer as TokTrait;
6use crate::error::{LayaError, Result};
7use std::path::Path;
8use tokenizers::Tokenizer as HfTokenizer;
9
10/// Fallback special ids of the Laya ModernBERT checkpoint (HF added tokens).
11const FB_CLS: u32 = 50281;
12const FB_SEP: u32 = 50282;
13const FB_MASK: u32 = 50284;
14const FB_PAD: u32 = 50283;
15const FB_UNK: u32 = 50280;
16
17/// Alias lookup order per special (ModernBERT `[X]` names first, Gemma `<x>` names after).
18const CLS_ALIASES: &[&str] = &["[CLS]", "<bos>", "<s>"];
19const SEP_ALIASES: &[&str] = &["[SEP]", "<eos>", "</s>"];
20const PAD_ALIASES: &[&str] = &["[PAD]", "<pad>"];
21const MASK_ALIASES: &[&str] = &["[MASK]", "<mask>"];
22const UNK_ALIASES: &[&str] = &["[UNK]", "<unk>"];
23
24/// A loaded checkpoint tokenizer.
25pub struct LayaTokenizer {
26    inner: HfTokenizer,
27    cls: u32,
28    sep: u32,
29    mask: u32,
30    pad: u32,
31    #[allow(dead_code)]
32    unk: u32,
33    mask_token: String,
34}
35
36impl LayaTokenizer {
37    /// Load a tokenizer from a `tokenizer.json` file.
38    pub fn from_file<P: AsRef<Path>>(path: P) -> Result<Self> {
39        let inner = HfTokenizer::from_file(path.as_ref())
40            .map_err(|e| LayaError::Tokenizer(e.to_string()))?;
41        Ok(Self::from_hf(inner))
42    }
43
44    /// Encode with special tokens (CLS/SEP), truncated to `max_length`. Used for encoder
45    /// embeddings (mean-pooling), matching HF `tok(text, truncation=True, max_length=...)`.
46    pub fn encode_with_special(&self, text: &str, max_length: usize) -> Vec<u32> {
47        match self.inner.encode(text, true) {
48            Ok(enc) => {
49                let mut ids = enc.get_ids().to_vec();
50                ids.truncate(max_length);
51                ids
52            }
53            Err(_) => Vec::new(),
54        }
55    }
56
57    fn from_hf(inner: HfTokenizer) -> Self {
58        let (cls, _) = resolve(&inner, CLS_ALIASES, FB_CLS);
59        let (sep, _) = resolve(&inner, SEP_ALIASES, FB_SEP);
60        let (pad, _) = resolve(&inner, PAD_ALIASES, FB_PAD);
61        let (mask, mask_token) = resolve(&inner, MASK_ALIASES, FB_MASK);
62        let (unk, _) = resolve(&inner, UNK_ALIASES, FB_UNK);
63        LayaTokenizer {
64            inner,
65            cls,
66            sep,
67            mask,
68            pad,
69            unk,
70            mask_token,
71        }
72    }
73}
74
75/// Resolve a special id + its token text by trying each alias in order (added tokens or vocab).
76fn resolve(tok: &HfTokenizer, aliases: &[&str], fallback: u32) -> (u32, String) {
77    for a in aliases {
78        if let Some(id) = tok.token_to_id(a) {
79            return (id, (*a).to_string());
80        }
81    }
82    (fallback, aliases[0].to_string())
83}
84
85impl TokTrait for LayaTokenizer {
86    fn cls_id(&self) -> u32 {
87        self.cls
88    }
89    fn sep_id(&self) -> u32 {
90        self.sep
91    }
92    fn mask_id(&self) -> u32 {
93        self.mask
94    }
95    fn pad_id(&self) -> u32 {
96        self.pad
97    }
98    fn mask_token(&self) -> &str {
99        &self.mask_token
100    }
101    fn encode(&self, text: &str) -> Vec<u32> {
102        // add_special_tokens = false, matching `tok(text, add_special_tokens=False)`.
103        match self.inner.encode(text, false) {
104            Ok(enc) => enc.get_ids().to_vec(),
105            Err(_) => Vec::new(),
106        }
107    }
108}
laya/tokenizer.rs

laya/
tokenizer.rs
