splintr 0.19.1

Fast Rust tokenizer (BPE + SentencePiece + WordPiece) with Python bindings
Documentation
//! Compile-time agent-token ids, one module per bundled vocabulary.
//!
//! The constants themselves are generated — `include!`d from
//! `agent_tokens_generated.rs`, which
//! `scripts/generate_agent_tokens.py --lang rust` writes from the same table
//! that produces the Python `*_AGENT_TOKENS` classes. That shared table is the
//! point: the two languages drifted for a whole release cycle when Rust's
//! constants were hand-written and Python's were generated, leaving Rust with
//! two vocabularies and Python with seven.
//!
//! This file is the hand-written half, and holds only the tests that pin the
//! generated ids against the tokenizers they name — so regenerating never
//! clobbers them.

include!("agent_tokens_generated.rs");

#[cfg(test)]
mod tests {
    use crate::core::pretrained::{base_vocab_size, from_pretrained, PretrainedVocab};
    use crate::core::Tokenize;

    /// One row per module: the vocabulary it names, and a sample of its
    /// constants spanning the block — the first slot, a middle one, the two
    /// vocabularies most often claim for themselves, and the last.
    #[allow(clippy::type_complexity)]
    fn samples() -> Vec<(&'static str, [(&'static str, u32); 6])> {
        macro_rules! row {
            ($vocab:literal, $m:ident) => {
                (
                    $vocab,
                    [
                        ("<|system|>", super::$m::SYSTEM),
                        ("<|im_start|>", super::$m::IM_START),
                        ("<|im_end|>", super::$m::IM_END),
                        ("<|think|>", super::$m::THINK),
                        ("<|pad|>", super::$m::PAD),
                        ("<|/summary|>", super::$m::SUMMARY_END),
                    ],
                )
            };
        }
        vec![
            row!("cl100k_base", cl100k_agent_tokens),
            row!("o200k_base", o200k_agent_tokens),
            row!("gpt-oss", gpt_oss_agent_tokens),
            row!("llama3", llama3_agent_tokens),
            row!("qwen3", qwen3_agent_tokens),
            row!("glm4", glm4_agent_tokens),
            row!("deepseek_v3", deepseek_v3_agent_tokens),
            row!("mistral_v1", mistral_v1_agent_tokens),
            row!("mistral_v2", mistral_v2_agent_tokens),
            row!("mistral_v3", mistral_v3_agent_tokens),
            row!("kimi_k2", kimi_k2_agent_tokens),
            row!("kimi_k3", kimi_k3_agent_tokens),
            row!("phi4", phi4_agent_tokens),
            row!("olmo2", olmo2_agent_tokens),
            row!("llama2", llama2_agent_tokens),
            row!("codellama", codellama_agent_tokens),
            row!("modernbert", modernbert_agent_tokens),
            row!("gemma2", gemma2_agent_tokens),
            row!("gemma3", gemma3_agent_tokens),
            row!("gemma4", gemma4_agent_tokens),
        ]
    }

    /// Every bundled vocabulary that carries agent tokens has a module here.
    ///
    /// The three Gemma vocabularies were bundled with working agent tokens and
    /// no module for a whole release: `insert_agent_tokens` gave each of them
    /// all 54, while the generator — which produces these constants, the Python
    /// `*_AGENT_TOKENS` classes and the tables in `docs/special_tokens.md` — had
    /// never heard of them. Nothing failed, because nothing asked. This asks.
    ///
    /// The vocabulary set comes from [`PretrainedVocab::supported_names`] rather
    /// than a list written here, because a list written here is exactly what was
    /// already silently incomplete. A family reachable by no name is a family no
    /// caller can load.
    #[test]
    fn every_vocabulary_with_agent_tokens_has_a_module() {
        let covered: Vec<PretrainedVocab> = samples()
            .into_iter()
            .map(|(name, _)| PretrainedVocab::from_name(name).expect("row names a vocabulary"))
            .collect();

        for name in PretrainedVocab::supported_names() {
            let vocab = PretrainedVocab::from_name(name).expect("a supported name resolves");
            if covered.contains(&vocab) {
                continue;
            }
            // Not covered: the only legitimate reason is carrying no agent
            // tokens at all, which is Whisper and is pinned just below.
            let tokenizer = from_pretrained(name).expect("bundled vocabulary loads");
            assert_eq!(
                tokenizer.special_token_id("<|think|>"),
                None,
                "{name} carries agent tokens but has no module in \
                 agent_tokens_generated.rs — add it to MODELS in \
                 scripts/generate_agent_tokens.py and regenerate"
            );
        }
    }

    /// Every generated constant must equal the id its tokenizer resolves that
    /// token to.
    ///
    /// This file and `insert_agent_tokens` in `pretrained.rs` are two
    /// independent statements of the same ids — this one generated by
    /// `scripts/generate_agent_tokens.py --lang rust`, that one written by
    /// hand. A constant is what callers reach for precisely so they need not
    /// think about the id, so a stale one is wrong in the worst way: silently.
    /// The Python half of the same contract is pinned by
    /// `python/tests/test_agent_token_constants.py`.
    #[test]
    fn generated_constants_match_the_tokenizers_they_name() {
        for (vocab, sample) in samples() {
            let tokenizer = from_pretrained(vocab).expect("bundled vocabulary loads");
            for (token, expected) in sample {
                assert_eq!(
                    tokenizer.special_token_id(token),
                    Some(expected),
                    "{vocab}: {token} disagrees with the vocabulary — regenerate with \
                     scripts/generate_agent_tokens.py --lang rust"
                );
            }
        }
    }

    /// Vocabularies that ship an agent-token name themselves keep their own id,
    /// below `base_vocab_size` — the id the checkpoint was trained on. A
    /// constant pointing at a splintr-appended id instead would build chat
    /// templates the model never saw.
    #[test]
    fn constants_defer_to_a_vocabularys_own_ids() {
        let qwen_base = base_vocab_size(PretrainedVocab::Qwen3);
        assert_eq!(super::qwen3_agent_tokens::IM_START, 151644);
        assert_eq!(super::qwen3_agent_tokens::IM_END, 151645);
        assert!(super::qwen3_agent_tokens::IM_START < qwen_base);
        // A name Qwen does not define still comes from the appended block.
        assert_eq!(super::qwen3_agent_tokens::SYSTEM, qwen_base);

        let glm_base = base_vocab_size(PretrainedVocab::Glm4);
        assert_eq!(super::glm4_agent_tokens::SYSTEM, 151335);
        assert_eq!(super::glm4_agent_tokens::IMAGE, 151363);
        assert!(super::glm4_agent_tokens::SYSTEM < glm_base);
        // GLM names only the opening markers, so the closing ones stay in the block.
        assert_eq!(super::glm4_agent_tokens::IMAGE_END, glm_base + 43);
    }

    /// Whisper is the one bundled vocabulary with no module here: it carries no
    /// agent tokens. Pinned so "no module" stays a stated fact rather than an
    /// omission someone later fills in by mistake.
    #[test]
    fn whisper_carries_no_agent_tokens() {
        let tokenizer = from_pretrained("whisper").expect("bundled vocabulary loads");
        assert_eq!(tokenizer.special_token_id("<|think|>"), None);
        assert_eq!(tokenizer.special_token_id("<|pad|>"), None);
        assert_eq!(
            base_vocab_size(PretrainedVocab::WhisperV2),
            tokenizer.vocab_size() as u32,
            "with no agent tokens, the base size is the whole vocabulary"
        );
    }
}