rlx-voxtral 0.2.9

// RLX — versatile ML compiler + runtime.
// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, version 3.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program. If not, see <https://www.gnu.org/licenses/>.

//! Host-side token embedding + audio placeholder fusion.

use crate::config::VoxtralConfig;
use crate::weights::VoxtralWeightPrefix;
use anyhow::{Result, bail, ensure};
use rlx_core::weight_map::WeightMap;
use std::path::Path;

const BOS_TOKEN_ID: u32 = 1;
const INST_TOKEN_ID: u32 = 3;
const BEGIN_AUDIO_TOKEN_ID: u32 = 25;
const END_INST_TOKEN_ID: u32 = 4;
const LANG_TOKEN_ID: u32 = 9909;
const COLON_TOKEN_ID: u32 = 1058;
const TRANSCRIBE_TOKEN_ID: u32 = 34;

/// Lookup token embeddings and scatter projected audio vectors at `audio_token_id` slots.
pub fn fuse_inputs_embeds(
    cfg: &VoxtralConfig,
    weights: &WeightMap,
    token_ids: &[u32],
    audio_embeds: &[f32],
) -> Result<Vec<f32>> {
    let h = cfg.text_config.hidden_size;
    let vocab = cfg.text_config.vocab_size;
    let embed_key = VoxtralWeightPrefix::lm_embed_tokens();
    let (embed, shape) = weights
        .get(embed_key)
        .ok_or_else(|| anyhow::anyhow!("missing {embed_key}"))?;
    ensure!(
        shape == [vocab, h],
        "unexpected embed shape {shape:?}, expected [{vocab}, {h}]"
    );
    let embed = embed.to_vec();

    let seq = token_ids.len();
    let n_audio_slots = token_ids
        .iter()
        .filter(|&&id| id == cfg.audio_token_id)
        .count();
    let n_audio_vecs = audio_embeds.len() / h;
    ensure!(
        n_audio_slots == n_audio_vecs,
        "audio token placeholders ({n_audio_slots}) != audio vectors ({n_audio_vecs})"
    );

    let mut out = vec![0f32; seq * h];
    let mut audio_idx = 0usize;
    for (pos, &tok) in token_ids.iter().enumerate() {
        if tok == cfg.audio_token_id {
            let src = &audio_embeds[audio_idx * h..(audio_idx + 1) * h];
            out[pos * h..(pos + 1) * h].copy_from_slice(src);
            audio_idx += 1;
            continue;
        }
        ensure!(
            (tok as usize) < vocab,
            "token id {tok} out of vocab range {vocab}"
        );
        let row = &embed[tok as usize * h..(tok as usize + 1) * h];
        out[pos * h..(pos + 1) * h].copy_from_slice(row);
    }
    Ok(out)
}

/// Greedy argmax on a logits row `[vocab]`.
pub fn argmax_token(logits: &[f32]) -> u32 {
    logits
        .iter()
        .enumerate()
        .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
        .map(|(i, _)| i as u32)
        .unwrap_or(0)
}

fn language_code_token_id(_model_dir: Option<&Path>, lang: &str) -> Result<u32> {
    match lang {
        "en" => Ok(1262),
        "fr" => Ok(7064),
        "de" => Ok(1558),
        "es" => Ok(4613),
        "pt" => Ok(8551),
        "hi" => Ok(6797),
        "nl" => Ok(7371),
        "it" => Ok(6360),
        other => bail!(
            "unknown language code {other:?}; place tekken.json next to weights or use en/fr/de/es/pt/hi/nl/it"
        ),
    }
}

/// Build a Mistral transcription prompt matching `apply_transcription_request`.
pub fn transcription_prompt_ids(
    cfg: &VoxtralConfig,
    n_audio: usize,
    language: Option<&str>,
    model_dir: Option<&Path>,
) -> Result<Vec<u32>> {
    let mut out = Vec::with_capacity(3 + n_audio + 6);
    out.extend([BOS_TOKEN_ID, INST_TOKEN_ID, BEGIN_AUDIO_TOKEN_ID]);
    out.extend(std::iter::repeat_n(cfg.audio_token_id, n_audio));
    out.push(END_INST_TOKEN_ID);
    if let Some(lang) = language {
        out.push(LANG_TOKEN_ID);
        out.push(COLON_TOKEN_ID);
        out.push(language_code_token_id(model_dir, lang)?);
    }
    out.push(TRANSCRIBE_TOKEN_ID);
    Ok(out)
}

pub fn decode_token_ids(model_dir: Option<&Path>, ids: &[u32]) -> Result<String> {
    #[cfg(feature = "tokenizer")]
    if let Some(dir) = model_dir {
        let tekken = dir.join("tekken.json");
        if tekken.is_file() {
            return tekken_decode(&tekken, ids);
        }
    }
    Ok(format!("{ids:?}"))
}

/// Decode token ids with a Mistral *tekken* (v7) tokenizer file.
///
/// HuggingFace `tokenizers` cannot parse tekken's schema, so this reads it
/// natively: ids `< num_special_tokens` are control/special tokens (dropped
/// from the rendered text, matching mistral-common's IGNORE policy); higher
/// ids index `vocab[id - num_special_tokens].token_bytes` (base64 raw bytes),
/// which are concatenated and rendered as UTF-8.
#[cfg(feature = "tokenizer")]
fn tekken_decode(tekken: &Path, ids: &[u32]) -> Result<String> {
    use anyhow::Context;
    let raw = std::fs::read_to_string(tekken).with_context(|| format!("read {tekken:?}"))?;
    let v: serde_json::Value =
        serde_json::from_str(&raw).with_context(|| format!("parse {tekken:?}"))?;
    let num_special = v["config"]["default_num_special_tokens"]
        .as_u64()
        .unwrap_or(1000) as usize;
    let vocab = v["vocab"]
        .as_array()
        .ok_or_else(|| anyhow::anyhow!("tekken: missing vocab array"))?;

    let mut bytes: Vec<u8> = Vec::new();
    for &id in ids {
        let id = id as usize;
        if id < num_special {
            continue; // control/special token — not part of rendered text
        }
        let entry = vocab
            .get(id - num_special)
            .ok_or_else(|| anyhow::anyhow!("tekken: token id {id} out of vocab range"))?;
        let b64 = entry["token_bytes"]
            .as_str()
            .ok_or_else(|| anyhow::anyhow!("tekken: token_bytes missing for id {id}"))?;
        base64_decode_into(b64, &mut bytes)?;
    }
    Ok(String::from_utf8_lossy(&bytes).into_owned())
}

/// Decode standard (RFC 4648) base64 into `out`. Stops at padding (`=`).
#[cfg(feature = "tokenizer")]
fn base64_decode_into(s: &str, out: &mut Vec<u8>) -> Result<()> {
    fn sextet(c: u8) -> Option<u32> {
        match c {
            b'A'..=b'Z' => Some((c - b'A') as u32),
            b'a'..=b'z' => Some((c - b'a' + 26) as u32),
            b'0'..=b'9' => Some((c - b'0' + 52) as u32),
            b'+' => Some(62),
            b'/' => Some(63),
            _ => None,
        }
    }
    let mut buf = 0u32;
    let mut bits = 0u32;
    for &c in s.as_bytes() {
        if c == b'=' {
            break;
        }
        let v = sextet(c).ok_or_else(|| anyhow::anyhow!("invalid base64 char {c:?}"))?;
        buf = (buf << 6) | v;
        bits += 6;
        if bits >= 8 {
            bits -= 8;
            out.push((buf >> bits) as u8);
        }
    }
    Ok(())
}

pub fn validate_prompt_audio_match(
    cfg: &VoxtralConfig,
    prompt: &[u32],
    n_audio: usize,
) -> Result<()> {
    let placeholders = prompt
        .iter()
        .filter(|&&id| id == cfg.audio_token_id)
        .count();
    if placeholders != n_audio {
        bail!(
            "prompt has {placeholders} audio placeholders (token {}), need {n_audio}",
            cfg.audio_token_id
        );
    }
    Ok(())
}

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

    #[test]
    fn transcription_prompt_matches_hf_template() {
        let cfg = VoxtralConfig::tiny_synthetic();
        let prompt = transcription_prompt_ids(&cfg, 4, Some("en"), None).unwrap();
        assert_eq!(
            prompt,
            vec![1, 3, 25, 24, 24, 24, 24, 4, 9909, 1058, 1262, 34]
        );
        let auto = transcription_prompt_ids(&cfg, 4, None, None).unwrap();
        assert_eq!(auto, vec![1, 3, 25, 24, 24, 24, 24, 4, 34]);
    }
}