//! Llama-3.2 BPE tokenizer wrapper.
//!
//! CSM speaker is encoded inline as text "[<id>]" then concatenated with the
//! utterance, no special speaker token id.

use crate::error::Result;
use std::path::Path;
use tokenizers::Tokenizer;

pub struct CsmTokenizer {
    inner: Tokenizer,
}

impl CsmTokenizer {
    pub fn from_file<P: AsRef<Path>>(path: P) -> Result<Self> {
        let inner = Tokenizer::from_file(path.as_ref())?;
        Ok(Self { inner })
    }

    pub fn format_segment(speaker: u32, text: &str) -> String {
        format!("[{speaker}]{text}")
    }

    pub fn encode(&self, text: &str) -> Result<Vec<u32>> {
        let enc = self.inner.encode(text, false)?;
        Ok(enc.get_ids().to_vec())
    }

    pub fn decode(&self, ids: &[u32]) -> Result<String> {
        Ok(self.inner.decode(ids, false)?)
    }
}