rustytorch/crates/models/rtx-csm/src/prompt.rs

//! Prompt assembly: convert `Segment`s into the (B, S, 33) token tensor + mask
//! that CSM's dual-transformer expects.
//!
//! Slot layout per row of the (S, cb+1) tensor:
//!
//!   - slots [0..cb):  per-codebook audio token id (or 0 if unused)
//!   - slot   cb:      Llama text token id (or 0 if unused)
//!
//! The mask (S, cb+1) u8 says which slots actually carry data.
//!
//! We delegate per-step tensor construction to candle's
//! `csm::Model::{audio_tokens_and_mask, text_tokens_and_mask}` and only do
//! the speaker formatting + concatenation here.

use crate::error::Result;
use crate::mimi::Mimi;
use crate::model::CsmModel;
use crate::tokenizer::CsmTokenizer;
use candle_core::Tensor;

#[derive(Debug, Clone)]
pub struct Segment {
    pub speaker: u32,
    pub text: String,
    /// 24 kHz mono f32; `None` means "to be generated".
    pub audio: Option<Vec<f32>>,
}

impl Segment {
    pub fn new_text(speaker: u32, text: impl Into<String>) -> Self {
        Self {
            speaker,
            text: text.into(),
            audio: None,
        }
    }

    pub fn new(speaker: u32, text: impl Into<String>, audio: Vec<f32>) -> Self {
        Self {
            speaker,
            text: text.into(),
            audio: Some(audio),
        }
    }
}

#[derive(Debug)]
pub struct PromptTensors {
    /// `(1, S, cb+1)` i64
    pub tokens: Tensor,
    /// `(1, S, cb+1)` u8
    pub mask: Tensor,
}

/// Tokenize a single segment into (tokens, mask) pieces. Audio encoding is
/// skipped when `segment.audio` is `None`.
pub fn encode_segment(
    segment: &Segment,
    model: &CsmModel,
    mimi: &mut Mimi,
    tokenizer: &CsmTokenizer,
) -> Result<PromptTensors> {
    // Text side: "[<speaker>]<text>" then tokenize with Llama BPE.
    let formatted = CsmTokenizer::format_segment(segment.speaker, &segment.text);
    let text_ids = tokenizer.encode(&formatted)?;
    let (text_tokens, text_mask) = model.inner.text_tokens_and_mask(&text_ids)?;

    // Audio side: run Mimi to get (1, cb, T) codes, transpose to (T, cb), then
    // emit one "audio frame" row per codec frame via candle's helper.
    let (audio_tokens, audio_mask) = if let Some(audio) = &segment.audio {
        let codes = mimi.encode(audio)?;
        let (_b, cb, t) = codes.dims3()?;
        assert_eq!(cb, model.config.audio_num_codebooks);

        let mut frame_tokens = Vec::with_capacity(t);
        let mut frame_masks = Vec::with_capacity(t);
        for frame_idx in 0..t {
            let frame = codes
                .narrow(2, frame_idx, 1)?
                .squeeze(2)?
                .flatten_all()?
                .to_vec1::<u32>()?;
            let (tok, mk) = model.inner.audio_tokens_and_mask(frame)?;
            frame_tokens.push(tok);
            frame_masks.push(mk);
        }
        if frame_tokens.is_empty() {
            (empty_tokens(model)?, empty_tokens(model)?)
        } else {
            let t = Tensor::cat(&frame_tokens, 1)?;
            let m = Tensor::cat(&frame_masks, 1)?;
            (t, m)
        }
    } else {
        (empty_tokens(model)?, empty_tokens(model)?)
    };

    // Concatenate text rows then audio rows along the sequence axis.
    let tokens = if audio_tokens.dim(1)? > 0 {
        Tensor::cat(&[&text_tokens, &audio_tokens], 1)?
    } else {
        text_tokens
    };
    let mask = if audio_mask.dim(1)? > 0 {
        Tensor::cat(&[&text_mask, &audio_mask], 1)?
    } else {
        text_mask
    };
    Ok(PromptTensors { tokens, mask })
}

fn empty_tokens(model: &CsmModel) -> Result<Tensor> {
    let cb = model.config.audio_num_codebooks;
    Ok(Tensor::zeros(
        (1, 0, cb + 1),
        candle_core::DType::U32,
        &model.device,
    )?)
}

/// Build the full prompt tensor = concat(context segments ..., current segment).
/// `current` is the utterance we want the model to continue from.
pub fn build_prompt(
    context: &[Segment],
    current: &Segment,
    model: &CsmModel,
    mimi: &mut Mimi,
    tokenizer: &CsmTokenizer,
) -> Result<PromptTensors> {
    let mut pieces_t: Vec<Tensor> = Vec::new();
    let mut pieces_m: Vec<Tensor> = Vec::new();
    for seg in context {
        let p = encode_segment(seg, model, mimi, tokenizer)?;
        pieces_t.push(p.tokens);
        pieces_m.push(p.mask);
    }
    let cur = encode_segment(current, model, mimi, tokenizer)?;
    pieces_t.push(cur.tokens);
    pieces_m.push(cur.mask);

    let tokens = Tensor::cat(&pieces_t.iter().collect::<Vec<_>>(), 1)?;
    let mask = Tensor::cat(&pieces_m.iter().collect::<Vec<_>>(), 1)?;
    Ok(PromptTensors { tokens, mask })
}