mozu-audio 0.3.0

//! Token-level contextual biasing with guarded acoustic tokens and written forms.

use crate::{
    Cancellation, Error, HaasEngine, InputLimits, Result, Tokenizer,
    result::WindowInference,
    window::{TimedToken, WindowRecognizer},
};
use sentencepiece_rs::SentencePieceProcessor;
use std::{
    collections::{BTreeMap, BTreeSet, VecDeque},
    path::{Path, PathBuf},
    sync::Arc,
    time::Instant,
};
use uuid::Uuid;

#[derive(Clone, Debug)]
pub struct VocabularyTerm {
    written: String,
    readings: Vec<String>,
}
impl VocabularyTerm {
    pub fn new(written: &str, readings: Vec<String>) -> Result<Self> {
        validate_form(written)?;
        for reading in &readings {
            validate_form(reading)?;
        }
        Ok(Self {
            written: written.to_owned(),
            readings,
        })
    }
    pub fn written(&self) -> &str {
        &self.written
    }
    pub fn readings(&self) -> &[String] {
        &self.readings
    }
}

#[derive(Clone, Copy, Debug)]
pub struct VocabularyConfig {
    /// Additive context weight on label logits; zero keeps acoustic decisions.
    pub strength: f32,
    /// At least one30ms frame must reach this RMS before biasing is applied.
    pub minimum_speech_rms: f32,
}
impl VocabularyConfig {
    pub const fn standard() -> Self {
        Self {
            strength: 2.0,
            minimum_speech_rms: 0.012,
        }
    }
}
#[derive(Clone, Debug)]
pub struct Vocabulary {
    terms: Vec<VocabularyTerm>,
    config: VocabularyConfig,
}
impl Vocabulary {
    pub fn new(terms: Vec<VocabularyTerm>, config: VocabularyConfig) -> Result<Self> {
        if terms.is_empty()
            || !config.strength.is_finite()
            || config.strength < 0.0
            || !config.minimum_speech_rms.is_finite()
            || config.minimum_speech_rms < 0.0
        {
            return Err(Error::Invalid(format!(
                "Invalid vocabulary: terms={}, config={config:?}",
                terms.len()
            )));
        }
        let mut written = BTreeSet::new();
        for term in &terms {
            if !written.insert(term.written.as_str()) {
                return Err(Error::Invalid(format!(
                    "Duplicate vocabulary written form: text={:?}",
                    term.written
                )));
            }
        }
        Ok(Self { terms, config })
    }
    pub fn terms(&self) -> &[VocabularyTerm] {
        &self.terms
    }
    pub fn config(&self) -> VocabularyConfig {
        self.config
    }
}

/// An owned request can cross to a model-owning service worker for compilation.
#[derive(Clone, Debug)]
pub struct VocabularyRequest {
    pub vocabulary: Vocabulary,
    pub tokenizer_model: PathBuf,
}

fn validate_form(text: &str) -> Result<()> {
    if text.is_empty() || text.trim() != text || text.chars().any(char::is_control) {
        return Err(Error::Invalid(format!(
            "Invalid vocabulary form: text={text:?}"
        )));
    }
    Ok(())
}
fn spellings(term: &VocabularyTerm) -> BTreeSet<String> {
    let mut forms = BTreeSet::new();
    for text in std::iter::once(&term.written).chain(&term.readings) {
        forms.insert(text.clone());
        forms.insert(text.to_lowercase());
        if text
            .chars()
            .all(|ch| ('\u{3040}'..='\u{30ff}').contains(&ch))
        {
            let katakana: String = text
                .chars()
                .map(|ch| {
                    if ('\u{3041}'..='\u{3096}').contains(&ch) {
                        char::from_u32(ch as u32 + 0x60).expect("hiragana maps to katakana")
                    } else {
                        ch
                    }
                })
                .collect();
            let hiragana: String = text
                .chars()
                .map(|ch| {
                    if ('\u{30a1}'..='\u{30f6}').contains(&ch) {
                        char::from_u32(ch as u32 - 0x60).expect("katakana maps to hiragana")
                    } else {
                        ch
                    }
                })
                .collect();
            forms.insert(katakana);
            forms.insert(hiragana);
        }
    }
    forms
}

fn tokenizer_error(path: &Path, operation: &str, error: sentencepiece_rs::Error) -> Error {
    Error::Invalid(format!(
        "Vocabulary tokenizer failed: path={}, operation={operation}, error={error}",
        path.display()
    ))
}

/// Remove only the utterance-initial marker. The remainder is represented by
/// existing pieces, with the model's byte pieces used for out-of-alphabet text.
fn inside_sequence(processor: &SentencePieceProcessor, ids: &[usize]) -> Result<Vec<usize>> {
    let Some(first) = ids.first() else {
        return Ok(Vec::new());
    };
    let piece = processor.id_to_piece(*first).map_err(|error| {
        Error::Invalid(format!(
            "Vocabulary first token failed: id={first}, error={error}"
        ))
    })?;
    let Some(remainder) = piece.strip_prefix('▁') else {
        return Ok(ids.to_vec());
    };
    let mut prefix = Vec::new();
    if !remainder.is_empty() {
        if let Some(id) = processor.model().try_piece_to_id(remainder) {
            prefix.push(id);
        } else {
            for ch in remainder.chars() {
                if let Some(id) = processor.model().try_piece_to_id(&ch.to_string()) {
                    prefix.push(id);
                } else {
                    let mut bytes = [0; 4];
                    for byte in ch.encode_utf8(&mut bytes).bytes() {
                        let name = format!("<0x{byte:02X}>");
                        prefix.push(processor.model().try_piece_to_id(&name).ok_or_else(|| Error::Invalid(format!("Vocabulary text is not representable inside an utterance: character={ch:?}, missing_piece={name}")))?);
                    }
                }
            }
        }
    }
    prefix.extend_from_slice(&ids[1..]);
    Ok(prefix)
}

#[derive(Clone, Debug)]
struct Pattern {
    tokens: Vec<usize>,
    term: usize,
}
#[derive(Clone, Debug)]
struct Node {
    edges: BTreeMap<usize, usize>,
    failure: usize,
    depth: usize,
    outputs: Vec<usize>,
    potential: f32,
}
impl Node {
    fn new(depth: usize) -> Self {
        Self {
            edges: BTreeMap::new(),
            failure: 0,
            depth,
            outputs: Vec::new(),
            potential: 0.0,
        }
    }
}
#[derive(Clone, Debug)]
struct FormatForm {
    surface: String,
    written: String,
}

/// Bound to the engine that validated the SentencePiece IDs. Compilation does
/// not change model weights. Written text can differ from the raw token pieces.
#[derive(Clone, Debug)]
pub struct CompiledVocabulary {
    engine_id: Uuid,
    vocabulary: Vocabulary,
    patterns: Vec<Pattern>,
    nodes: Vec<Node>,
    forms: Vec<FormatForm>,
}
impl CompiledVocabulary {
    pub(crate) fn compile(
        vocabulary: Vocabulary,
        tokenizer_model: &Path,
        tokenizer: &Tokenizer,
        engine_id: Uuid,
    ) -> Result<Self> {
        let processor = SentencePieceProcessor::open(tokenizer_model)
            .map_err(|error| tokenizer_error(tokenizer_model, "load", error))?;
        if processor.model().vocab_size() != tokenizer.pieces.len() {
            return Err(Error::Invalid(format!(
                "Vocabulary tokenizer size mismatch: path={}, expected={}, actual={}",
                tokenizer_model.display(),
                tokenizer.pieces.len(),
                processor.model().vocab_size()
            )));
        }
        for (id, expected) in tokenizer.pieces.iter().enumerate() {
            let actual = processor
                .id_to_piece(id)
                .map_err(|error| tokenizer_error(tokenizer_model, "verify token IDs", error))?;
            if actual != expected.text {
                return Err(Error::Invalid(format!(
                    "Vocabulary tokenizer ID mismatch: path={}, id={id}, expected={:?}, actual={actual:?}",
                    tokenizer_model.display(),
                    expected.text
                )));
            }
        }
        let mut patterns = Vec::new();
        let mut sequences: BTreeMap<Vec<usize>, usize> = BTreeMap::new();
        let mut forms = BTreeMap::new();
        for (term_index, term) in vocabulary.terms.iter().enumerate() {
            for spelling in spellings(term) {
                let ids = processor.encode_to_ids(&spelling).map_err(|error| {
                    tokenizer_error(tokenizer_model, "encode vocabulary", error)
                })?;
                if ids.is_empty() || ids.contains(&processor.unk_id()) {
                    return Err(Error::Invalid(format!(
                        "Vocabulary contains an empty or unknown encoding: spelling={spelling:?}, ids={ids:?}"
                    )));
                }
                let surface = processor.decode_ids(&ids).map_err(|error| {
                    tokenizer_error(tokenizer_model, "normalize vocabulary form", error)
                })?;
                if let Some(previous) = forms.insert(surface.clone(), term.written.clone())
                    && previous != term.written
                {
                    return Err(Error::Invalid(format!(
                        "Ambiguous vocabulary reading: reading={surface:?}, first={previous:?}, second={:?}",
                        term.written
                    )));
                }
                let inside = inside_sequence(&processor, &ids)?;
                for sequence in [ids, inside] {
                    if sequence.is_empty() {
                        continue;
                    }
                    if let Some(previous) = sequences.insert(sequence.clone(), term_index) {
                        if previous != term_index {
                            return Err(Error::Invalid(format!(
                                "Vocabulary token sequence has conflicting written forms: ids={sequence:?}"
                            )));
                        }
                        continue;
                    }
                    patterns.push(Pattern {
                        tokens: sequence,
                        term: term_index,
                    });
                }
            }
        }
        let nodes = build_graph(&patterns);
        let mut forms: Vec<_> = forms
            .into_iter()
            .map(|(surface, written)| FormatForm { surface, written })
            .collect();
        forms.sort_by_key(|form| std::cmp::Reverse(form.surface.len()));
        Ok(Self {
            engine_id,
            vocabulary,
            patterns,
            nodes,
            forms,
        })
    }
    pub(crate) fn validate_engine(&self, engine_id: Uuid) -> Result<()> {
        if self.engine_id != engine_id {
            return Err(Error::Invalid(format!(
                "Vocabulary belongs to a different loaded engine: compiled={}, requested={engine_id}",
                self.engine_id
            )));
        }
        Ok(())
    }
    pub fn vocabulary(&self) -> &Vocabulary {
        &self.vocabulary
    }
    pub fn token_sequences(&self) -> Vec<Vec<usize>> {
        self.patterns
            .iter()
            .map(|pattern| pattern.tokens.clone())
            .collect()
    }
    pub(crate) fn has_speech(&self, audio: &[f32]) -> bool {
        audio.chunks(480).any(|frame| {
            let power = frame
                .iter()
                .map(|sample| f64::from(*sample).powi(2))
                .sum::<f64>();
            power > 0.0
                && (power / frame.len() as f64).sqrt()
                    >= self.vocabulary.config.minimum_speech_rms as f64
        })
    }
    fn next(&self, state: usize, token: usize) -> usize {
        transition(&self.nodes, state, token)
    }
    pub(crate) fn format(&self, text: &str) -> String {
        let mut result = String::new();
        let mut offset = 0;
        while offset < text.len() {
            let found = self.forms.iter().find(|form| {
                let end = offset + form.surface.len();
                text[offset..].starts_with(&form.surface)
                    && !(form.surface.as_bytes()[0].is_ascii_alphanumeric()
                        && offset > 0
                        && text.as_bytes()[offset - 1].is_ascii_alphanumeric())
                    && !(form.surface.as_bytes()[form.surface.len() - 1].is_ascii_alphanumeric()
                        && end < text.len()
                        && text.as_bytes()[end].is_ascii_alphanumeric())
            });
            if let Some(form) = found {
                result.push_str(&form.written);
                offset += form.surface.len();
            } else {
                let ch = text[offset..]
                    .chars()
                    .next()
                    .expect("offset is a character boundary");
                result.push(ch);
                offset += ch.len_utf8();
            }
        }
        result
    }
    fn completed_ranges(&self, tokens: &[TimedToken]) -> Vec<std::ops::Range<usize>> {
        let mut state = 0;
        let mut ranges = Vec::new();
        for (index, token) in tokens.iter().enumerate() {
            state = self.next(state, token.token_id);
            for output in &self.nodes[state].outputs {
                let pattern = &self.patterns[*output];
                let _term = &self.vocabulary.terms[pattern.term];
                ranges.push(index + 1 - pattern.tokens.len()..index + 1);
            }
        }
        ranges
    }
    pub(crate) fn stable_prefix(&self, tokens: &[TimedToken], proposed: usize) -> usize {
        let mut state = 0;
        for token in &tokens[..proposed] {
            state = self.next(state, token.token_id);
        }
        proposed.saturating_sub(self.nodes[state].potential as usize)
    }
    /// Keep acoustic tokens outside complete vocabulary replacements. Partial
    /// word bonuses never authorize unrelated insertions or deleted neighbours.
    pub(crate) fn guard(&self, plain: &[TimedToken], boosted: &[TimedToken]) -> Vec<TimedToken> {
        let ranges = self.completed_ranges(boosted);
        if ranges.is_empty() {
            return plain.to_vec();
        }
        let width = boosted.len() + 1;
        let mut lcs = vec![0usize; (plain.len() + 1) * width];
        for i in (0..plain.len()).rev() {
            for j in (0..boosted.len()).rev() {
                lcs[i * width + j] = if plain[i].token_id == boosted[j].token_id {
                    lcs[(i + 1) * width + j + 1] + 1
                } else {
                    lcs[(i + 1) * width + j].max(lcs[i * width + j + 1])
                };
            }
        }
        let mut anchors = Vec::new();
        let (mut i, mut j) = (0, 0);
        while i < plain.len() && j < boosted.len() {
            if plain[i].token_id == boosted[j].token_id {
                anchors.push((i, j));
                i += 1;
                j += 1;
            } else if lcs[(i + 1) * width + j] >= lcs[i * width + j + 1] {
                i += 1;
            } else {
                j += 1;
            }
        }
        anchors.push((plain.len(), boosted.len()));
        let mut output = Vec::new();
        let (mut left, mut right) = (0, 0);
        for (i, j) in anchors {
            let accepted = ranges
                .iter()
                .find(|range| right < j && range.start <= right && j <= range.end);
            if let Some(range) = accepted {
                let start = boosted[range.start].start_sample;
                let last = &boosted[range.end - 1];
                let end = last.start_sample.saturating_add(
                    last.duration_samples
                        .max(crate::window::ENCODER_STEP as u64),
                );
                output.extend(
                    plain[left..i]
                        .iter()
                        .filter(|token| token.start_sample < start)
                        .cloned(),
                );
                output.extend_from_slice(&boosted[right..j]);
                output.extend(
                    plain[left..i]
                        .iter()
                        .filter(|token| token.start_sample >= end)
                        .cloned(),
                );
            } else {
                output.extend_from_slice(&plain[left..i]);
            }
            if i < plain.len() {
                output.push(plain[i].clone());
            }
            left = i + 1;
            right = j + 1;
        }
        // A large duration jump can shift neighbouring acoustic tokens. Reject
        // a replacement that would break the public monotonic time contract.
        if output
            .windows(2)
            .any(|pair| pair[0].start_sample > pair[1].start_sample)
        {
            return plain.to_vec();
        }
        output
    }
}

pub(crate) struct VocabularyRecognizer<'a> {
    engine: &'a HaasEngine,
    vocabulary: Arc<CompiledVocabulary>,
}
impl<'a> VocabularyRecognizer<'a> {
    pub(crate) fn new(engine: &'a HaasEngine, vocabulary: Arc<CompiledVocabulary>) -> Result<Self> {
        vocabulary.validate_engine(engine.vocabulary_id)?;
        Ok(Self { engine, vocabulary })
    }
}
impl WindowRecognizer for VocabularyRecognizer<'_> {
    fn input_limits(&self) -> InputLimits {
        self.engine.input_limits()
    }
    fn infer(&self, audio: &[f32], cancellation: &Cancellation) -> Result<Vec<TimedToken>> {
        Ok(self
            .engine
            .transcribe_samples_with_vocabulary(audio, 16000, &self.vocabulary, cancellation)?
            .token_timings)
    }
    fn infer_with_metrics(
        &self,
        audio: &[f32],
        original_samples: usize,
        cancellation: &Cancellation,
    ) -> Result<WindowInference> {
        let start = Instant::now();
        let result = self.engine.transcribe_samples_with_vocabulary(
            audio,
            16000,
            &self.vocabulary,
            cancellation,
        )?;
        Ok(WindowInference::from_recognition(
            result,
            audio.len(),
            original_samples,
            start.elapsed().as_secs_f64(),
        ))
    }
    fn render(&self, tokens: &[TimedToken]) -> Result<String> {
        let raw = self.engine.assets.tokenizer.decode(
            &tokens
                .iter()
                .map(|token| token.token_id)
                .collect::<Vec<_>>(),
        )?;
        Ok(self.vocabulary.format(&raw))
    }
    fn token_piece(&self, id: usize) -> Result<crate::result::TokenPiece> {
        self.engine.token_piece(id)
    }
    fn stable_prefix(&self, tokens: &[TimedToken], proposed: usize) -> usize {
        self.vocabulary.stable_prefix(tokens, proposed)
    }
}

fn transition(nodes: &[Node], mut state: usize, token: usize) -> usize {
    loop {
        if let Some(next) = nodes[state].edges.get(&token) {
            return *next;
        }
        if state == 0 {
            return 0;
        }
        state = nodes[state].failure;
    }
}
fn build_graph(patterns: &[Pattern]) -> Vec<Node> {
    let mut nodes = vec![Node::new(0)];
    for (index, pattern) in patterns.iter().enumerate() {
        let mut state = 0;
        for token in &pattern.tokens {
            state = if let Some(next) = nodes[state].edges.get(token) {
                *next
            } else {
                let next = nodes.len();
                let depth = nodes[state].depth + 1;
                nodes.push(Node::new(depth));
                nodes[state].edges.insert(*token, next);
                next
            };
        }
        nodes[state].outputs.push(index);
    }
    let mut queue: VecDeque<_> = nodes[0].edges.values().copied().collect();
    while let Some(state) = queue.pop_front() {
        let failure = nodes[state].failure;
        let suffix_outputs = nodes[failure].outputs.clone();
        nodes[state].outputs.extend(suffix_outputs);
        nodes[state].potential = if nodes[state].edges.is_empty() {
            nodes[failure].potential
        } else {
            nodes[state].depth as f32
        };
        let edges: Vec<_> = nodes[state]
            .edges
            .iter()
            .map(|(token, next)| (*token, *next))
            .collect();
        for (token, next) in edges {
            nodes[next].failure = transition(&nodes, failure, token);
            queue.push_back(next);
        }
    }
    nodes
}

pub(crate) struct BiasState<'a> {
    vocabulary: &'a CompiledVocabulary,
    state: usize,
}
impl<'a> BiasState<'a> {
    pub(crate) fn new(vocabulary: &'a CompiledVocabulary) -> Self {
        Self {
            vocabulary,
            state: 0,
        }
    }
    pub(crate) fn select(&self, logits: &[f32], blank: usize) -> Result<usize> {
        if logits.is_empty()
            || blank >= logits.len()
            || logits.iter().any(|value| !value.is_finite())
        {
            return Err(Error::Invalid(format!(
                "Invalid vocabulary label logits: classes={}, blank={blank}",
                logits.len()
            )));
        }
        let strength = self.vocabulary.vocabulary.config.strength;
        let debt = self.vocabulary.nodes[self.state].potential * strength;
        let mut best = 0;
        let mut best_score = if best == blank {
            logits[0]
        } else {
            logits[0] - debt
        };
        for (token, value) in logits.iter().enumerate().skip(1) {
            let score = if token == blank {
                *value
            } else {
                *value - debt
            };
            if score > best_score {
                best = token;
                best_score = score;
            }
        }
        let mut candidates = BTreeSet::new();
        let mut state = self.state;
        loop {
            candidates.extend(self.vocabulary.nodes[state].edges.keys().copied());
            if state == 0 {
                break;
            }
            state = self.vocabulary.nodes[state].failure;
        }
        for token in candidates {
            if token == blank || token >= logits.len() {
                return Err(Error::Invalid(format!(
                    "Vocabulary contains invalid label: token={token}, blank={blank}, classes={}",
                    logits.len()
                )));
            }
            let next = self.vocabulary.next(self.state, token);
            let node = &self.vocabulary.nodes[next];
            let completion = node
                .outputs
                .iter()
                .map(|index| self.vocabulary.patterns[*index].tokens.len())
                .max()
                .unwrap_or(0) as f32;
            let score = logits[token] - debt + (node.potential + completion) * strength;
            if !score.is_finite() {
                return Err(Error::Invalid(format!(
                    "Vocabulary score overflow: token={token}, strength={strength}"
                )));
            }
            if score > best_score || score == best_score && token < best {
                best = token;
                best_score = score;
            }
        }
        Ok(best)
    }
    pub(crate) fn accept(&mut self, token: usize) {
        self.state = self.vocabulary.next(self.state, token);
    }
}

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

    fn context(sequences: Vec<Vec<usize>>, strength: f32) -> CompiledVocabulary {
        let vocabulary = Vocabulary::new(
            vec![VocabularyTerm::new("Haas", vec!["ハース".to_owned()]).unwrap()],
            VocabularyConfig {
                strength,
                minimum_speech_rms: 0.012,
            },
        )
        .unwrap();
        let patterns: Vec<_> = sequences
            .into_iter()
            .map(|tokens| Pattern { tokens, term: 0 })
            .collect();
        let nodes = build_graph(&patterns);
        CompiledVocabulary {
            engine_id: Uuid::nil(),
            vocabulary,
            patterns,
            nodes,
            forms: vec![
                FormatForm {
                    surface: "ハース".to_owned(),
                    written: "Haas".to_owned(),
                },
                FormatForm {
                    surface: "haas".to_owned(),
                    written: "Haas".to_owned(),
                },
            ],
        }
    }
    fn tokens(ids: &[usize]) -> Vec<TimedToken> {
        ids.iter()
            .enumerate()
            .map(|(index, id)| TimedToken {
                token_id: *id,
                start_sample: index as u64 * 1280,
                duration_samples: 1280,
                confidence: 0.8,
            })
            .collect()
    }
    #[test]
    fn vocabulary_changes_joint_selection_and_zero_weight_preserves_acoustics() {
        let vocabulary = context(vec![vec![1, 2]], 2.0);
        let mut state = BiasState::new(&vocabulary);
        assert_eq!(state.select(&[0.1, -0.5, -0.6, 0.0], 3).unwrap(), 1);
        state.accept(1);
        assert_eq!(state.select(&[3.0, -10.0, 0.0, -10.0], 3).unwrap(), 2);
        state.accept(2);
        assert_eq!(state.select(&[1.0, -10.0, -10.0, 0.0], 3).unwrap(), 0);
        let neutral = context(vec![vec![1, 2]], 0.0);
        let mut state = BiasState::new(&neutral);
        for prefix in [0, 1, 2] {
            state.accept(prefix);
            for values in [
                [1.0, 0.2, 0.3, 0.1],
                [0.0, 0.0, 0.0, 0.0],
                [-2.0, -1.0, -3.0, 0.0],
            ] {
                let expected = (1..values.len()).fold(0, |best, index| {
                    if values[index] > values[best] {
                        index
                    } else {
                        best
                    }
                });
                assert_eq!(state.select(&values, 3).unwrap(), expected);
            }
        }
    }
    #[test]
    fn incomplete_words_do_not_authorize_changes_and_silence_is_not_boosted() {
        let vocabulary = context(vec![vec![1, 2]], 2.0);
        let plain = tokens(&[0, 3, 4]);
        assert_eq!(vocabulary.guard(&plain, &tokens(&[0, 1, 4])), plain);
        assert!(!vocabulary.has_speech(&vec![0.0; 16000]));
        assert!(!vocabulary.has_speech(&vec![0.001; 16000]));
        assert!(vocabulary.has_speech(&vec![0.02; 480]));
    }
    #[test]
    fn guarded_word_replacement_preserves_a_neighbour_deleted_by_biased_decoding() {
        let vocabulary = context(vec![vec![1, 2]], 2.0);
        let plain = tokens(&[0, 3, 4, 5, 9]);
        let mut boosted = tokens(&[0, 1, 2, 9]);
        boosted[3].start_sample = 5120;
        let guarded = vocabulary.guard(&plain, &boosted);
        assert_eq!(
            guarded
                .iter()
                .map(|token| token.token_id)
                .collect::<Vec<_>>(),
            vec![0, 1, 2, 5, 9]
        );
        assert!(
            guarded
                .windows(2)
                .all(|pair| pair[0].start_sample <= pair[1].start_sample)
        );
    }
    #[test]
    fn written_forms_and_stability_do_not_rewrite_an_already_committed_prefix() {
        let vocabulary = context(vec![vec![1, 2], vec![1, 2, 3]], 2.0);
        assert_eq!(vocabulary.stable_prefix(&tokens(&[0, 1, 2]), 3), 1);
        assert_eq!(vocabulary.stable_prefix(&tokens(&[0, 1, 2, 3]), 4), 4);
        assert_eq!(
            vocabulary.format("ハースとhaas、haasx"),
            "HaasとHaas、haasx"
        );
    }
}