path: root/makima/src/tts/qwen3/model.rs

//! Qwen3 Language Model transformer backbone.
//!
//! Implements the 28-layer transformer with:
//! - Rotary Position Embeddings (RoPE)
//! - Grouped Query Attention (GQA) — 16 heads, 8 KV heads
//! - SiLU-gated MLP
//! - RMS normalization
//! - KV cache for autoregressive generation
//!
//! Based on the candle-transformers Qwen2 model architecture,
//! extended for Qwen3-TTS.

use candle_core::{DType, Device, IndexOp, Module, Result, Tensor, D};
use candle_nn::{embedding, linear_no_bias, rms_norm, Embedding, Linear, RmsNorm, VarBuilder};

use super::config::Qwen3LmConfig;

// ---------------------------------------------------------------------------
// Rotary Position Embeddings
// ---------------------------------------------------------------------------

/// Precomputed RoPE sin/cos tables.
#[derive(Debug, Clone)]
pub struct RotaryEmbedding {
    cos: Tensor,
    sin: Tensor,
}

impl RotaryEmbedding {
    pub fn new(config: &Qwen3LmConfig, dtype: DType, device: &Device) -> Result<Self> {
        let head_dim = config.head_dim;
        let max_seq = config.max_position_embeddings;
        let theta = config.rope_theta;

        let inv_freq: Vec<f32> = (0..head_dim)
            .step_by(2)
            .map(|i| 1.0 / (theta as f32).powf(i as f32 / head_dim as f32))
            .collect();

        let inv_freq_tensor =
            Tensor::from_vec(inv_freq, (head_dim / 2,), device)?.to_dtype(DType::F32)?;

        let positions: Vec<f32> = (0..max_seq).map(|p| p as f32).collect();
        let positions_tensor = Tensor::from_vec(positions, (max_seq, 1), device)?;

        // [max_seq, head_dim/2]
        let freqs = positions_tensor.matmul(&inv_freq_tensor.unsqueeze(0)?)?;
        // [max_seq, head_dim] by repeating
        let emb = Tensor::cat(&[&freqs, &freqs], D::Minus1)?;

        let cos = emb.cos()?.to_dtype(dtype)?;
        let sin = emb.sin()?.to_dtype(dtype)?;

        Ok(Self { cos, sin })
    }

    /// Apply RoPE to query and key tensors.
    /// Input shape: [batch, heads, seq_len, head_dim]
    pub fn apply(&self, q: &Tensor, k: &Tensor, offset: usize) -> Result<(Tensor, Tensor)> {
        let seq_len = q.dim(2)?;
        let cos = self.cos.narrow(0, offset, seq_len)?;
        let sin = self.sin.narrow(0, offset, seq_len)?;

        let cos = cos.unsqueeze(0)?.unsqueeze(0)?; // [1, 1, seq, dim]
        let sin = sin.unsqueeze(0)?.unsqueeze(0)?;

        let q_rotated = Self::rotate_half(q, &cos, &sin)?;
        let k_rotated = Self::rotate_half(k, &cos, &sin)?;

        Ok((q_rotated, k_rotated))
    }

    fn rotate_half(x: &Tensor, cos: &Tensor, sin: &Tensor) -> Result<Tensor> {
        let half_dim = x.dim(D::Minus1)? / 2;
        let x1 = x.narrow(D::Minus1, 0, half_dim)?;
        let x2 = x.narrow(D::Minus1, half_dim, half_dim)?;

        // [-x2, x1] concatenated
        let neg_x2 = x2.neg()?;
        let rotated = Tensor::cat(&[&neg_x2, &x1], D::Minus1)?;

        // x * cos + rotated * sin
        let result = x.broadcast_mul(cos)?.broadcast_add(&rotated.broadcast_mul(sin)?)?;
        Ok(result)
    }
}

// ---------------------------------------------------------------------------
// KV Cache
// ---------------------------------------------------------------------------

/// Per-layer key-value cache for autoregressive generation.
#[derive(Debug, Clone)]
pub struct KvCache {
    key: Option<Tensor>,
    value: Option<Tensor>,
}

impl KvCache {
    pub fn new() -> Self {
        Self {
            key: None,
            value: None,
        }
    }

    /// Append new key/value tensors and return the full cached sequence.
    /// Input shapes: [batch, num_kv_heads, new_seq_len, head_dim]
    pub fn append(&mut self, key: &Tensor, value: &Tensor) -> Result<(Tensor, Tensor)> {
        let (full_key, full_value) = match (&self.key, &self.value) {
            (Some(prev_k), Some(prev_v)) => {
                let k = Tensor::cat(&[prev_k, key], 2)?;
                let v = Tensor::cat(&[prev_v, value], 2)?;
                (k, v)
            }
            _ => (key.clone(), value.clone()),
        };

        self.key = Some(full_key.clone());
        self.value = Some(full_value.clone());

        Ok((full_key, full_value))
    }

    /// Current cached sequence length.
    pub fn seq_len(&self) -> usize {
        self.key
            .as_ref()
            .map(|k| k.dim(2).unwrap_or(0))
            .unwrap_or(0)
    }

    /// Reset the cache.
    pub fn reset(&mut self) {
        self.key = None;
        self.value = None;
    }
}

// ---------------------------------------------------------------------------
// Attention
// ---------------------------------------------------------------------------

/// Multi-head attention with GQA and RoPE.
pub struct Qwen3Attention {
    q_proj: Linear,
    k_proj: Linear,
    v_proj: Linear,
    o_proj: Linear,
    q_norm: RmsNorm,
    k_norm: RmsNorm,
    num_heads: usize,
    num_kv_heads: usize,
    head_dim: usize,
    num_kv_groups: usize,
}

impl Qwen3Attention {
    pub fn new(config: &Qwen3LmConfig, vb: VarBuilder) -> Result<Self> {
        let hidden = config.hidden_size;
        let num_heads = config.num_attention_heads;
        let num_kv_heads = config.num_key_value_heads;
        let head_dim = config.head_dim;

        let q_proj = linear_no_bias(hidden, num_heads * head_dim, vb.pp("q_proj"))?;
        let k_proj = linear_no_bias(hidden, num_kv_heads * head_dim, vb.pp("k_proj"))?;
        let v_proj = linear_no_bias(hidden, num_kv_heads * head_dim, vb.pp("v_proj"))?;
        let o_proj = linear_no_bias(num_heads * head_dim, hidden, vb.pp("o_proj"))?;

        let q_norm = rms_norm(head_dim, config.rms_norm_eps, vb.pp("q_norm"))?;
        let k_norm = rms_norm(head_dim, config.rms_norm_eps, vb.pp("k_norm"))?;

        Ok(Self {
            q_proj,
            k_proj,
            v_proj,
            o_proj,
            q_norm,
            k_norm,
            num_heads,
            num_kv_heads,
            head_dim,
            num_kv_groups: config.num_kv_groups(),
        })
    }

    /// Forward pass with KV cache and RoPE.
    /// Input: [batch, seq_len, hidden_size]
    /// Returns: [batch, seq_len, hidden_size]
    pub fn forward(
        &self,
        hidden_states: &Tensor,
        rope: &RotaryEmbedding,
        kv_cache: &mut KvCache,
        attention_mask: Option<&Tensor>,
    ) -> Result<Tensor> {
        let (batch, seq_len, _) = hidden_states.dims3()?;
        let offset = kv_cache.seq_len();

        // Project Q, K, V
        let q = self.q_proj.forward(hidden_states)?;
        let k = self.k_proj.forward(hidden_states)?;
        let v = self.v_proj.forward(hidden_states)?;

        // Reshape: [batch, seq, heads*dim] -> [batch, heads, seq, dim]
        let q = q
            .reshape((batch, seq_len, self.num_heads, self.head_dim))?
            .transpose(1, 2)?;
        let k = k
            .reshape((batch, seq_len, self.num_kv_heads, self.head_dim))?
            .transpose(1, 2)?;
        let v = v
            .reshape((batch, seq_len, self.num_kv_heads, self.head_dim))?
            .transpose(1, 2)?;

        // Apply QK normalization (Qwen3 specific)
        let q = self.apply_head_norm(&q, &self.q_norm)?;
        let k = self.apply_head_norm(&k, &self.k_norm)?;

        // Apply RoPE
        let (q, k) = rope.apply(&q, &k, offset)?;

        // Update KV cache
        let (k, v) = kv_cache.append(&k, &v)?;

        // Expand KV heads for GQA: [batch, kv_heads, seq, dim] -> [batch, heads, seq, dim]
        let k = self.repeat_kv(&k)?;
        let v = self.repeat_kv(&v)?;

        // Scaled dot-product attention
        let scale = (self.head_dim as f64).sqrt();
        let attn_weights = (q.matmul(&k.transpose(D::Minus2, D::Minus1)?)? / scale)?;

        let attn_weights = match attention_mask {
            Some(mask) => attn_weights.broadcast_add(mask)?,
            None => attn_weights,
        };

        let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?;

        // Attention output
        let attn_output = attn_weights.matmul(&v)?;

        // [batch, heads, seq, dim] -> [batch, seq, heads*dim]
        let attn_output = attn_output
            .transpose(1, 2)?
            .reshape((batch, seq_len, self.num_heads * self.head_dim))?;

        self.o_proj.forward(&attn_output)
    }

    /// Apply RMS norm per-head.
    fn apply_head_norm(&self, x: &Tensor, norm: &RmsNorm) -> Result<Tensor> {
        let (b, h, s, d) = x.dims4()?;
        // Reshape to [b*h*s, d] for norm, then back
        let flat = x.reshape((b * h * s, d))?;
        let normed = norm.forward(&flat)?;
        normed.reshape((b, h, s, d))
    }

    /// Repeat KV heads for GQA.
    fn repeat_kv(&self, x: &Tensor) -> Result<Tensor> {
        if self.num_kv_groups == 1 {
            return Ok(x.clone());
        }
        let (batch, num_kv_heads, seq_len, head_dim) = x.dims4()?;
        let x = x
            .unsqueeze(2)?
            .expand((batch, num_kv_heads, self.num_kv_groups, seq_len, head_dim))?
            .reshape((batch, self.num_heads, seq_len, head_dim))?;
        Ok(x)
    }
}

// ---------------------------------------------------------------------------
// MLP
// ---------------------------------------------------------------------------

/// SiLU-gated feed-forward network.
pub struct Qwen3Mlp {
    gate_proj: Linear,
    up_proj: Linear,
    down_proj: Linear,
}

impl Qwen3Mlp {
    pub fn new(config: &Qwen3LmConfig, vb: VarBuilder) -> Result<Self> {
        let hidden = config.hidden_size;
        let intermediate = config.intermediate_size;

        let gate_proj = linear_no_bias(hidden, intermediate, vb.pp("gate_proj"))?;
        let up_proj = linear_no_bias(hidden, intermediate, vb.pp("up_proj"))?;
        let down_proj = linear_no_bias(intermediate, hidden, vb.pp("down_proj"))?;

        Ok(Self {
            gate_proj,
            up_proj,
            down_proj,
        })
    }

    pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
        let gate = self.gate_proj.forward(x)?;
        let gate = candle_nn::Activation::Silu.forward(&gate)?;
        let up = self.up_proj.forward(x)?;
        let hidden = (gate * up)?;
        self.down_proj.forward(&hidden)
    }
}

// ---------------------------------------------------------------------------
// Transformer Layer
// ---------------------------------------------------------------------------

/// A single Qwen3 transformer decoder layer.
pub struct Qwen3DecoderLayer {
    self_attn: Qwen3Attention,
    mlp: Qwen3Mlp,
    input_layernorm: RmsNorm,
    post_attention_layernorm: RmsNorm,
}

impl Qwen3DecoderLayer {
    pub fn new(config: &Qwen3LmConfig, vb: VarBuilder) -> Result<Self> {
        let self_attn = Qwen3Attention::new(config, vb.pp("self_attn"))?;
        let mlp = Qwen3Mlp::new(config, vb.pp("mlp"))?;
        let input_layernorm =
            rms_norm(config.hidden_size, config.rms_norm_eps, vb.pp("input_layernorm"))?;
        let post_attention_layernorm = rms_norm(
            config.hidden_size,
            config.rms_norm_eps,
            vb.pp("post_attention_layernorm"),
        )?;

        Ok(Self {
            self_attn,
            mlp,
            input_layernorm,
            post_attention_layernorm,
        })
    }

    pub fn forward(
        &self,
        hidden_states: &Tensor,
        rope: &RotaryEmbedding,
        kv_cache: &mut KvCache,
        attention_mask: Option<&Tensor>,
    ) -> Result<Tensor> {
        // Pre-norm attention
        let residual = hidden_states;
        let hidden_states = self.input_layernorm.forward(hidden_states)?;
        let hidden_states =
            self.self_attn
                .forward(&hidden_states, rope, kv_cache, attention_mask)?;
        let hidden_states = (residual + hidden_states)?;

        // Pre-norm MLP
        let residual = &hidden_states;
        let hidden_states = self.post_attention_layernorm.forward(&hidden_states)?;
        let hidden_states = self.mlp.forward(&hidden_states)?;
        let output = (residual + hidden_states)?;

        Ok(output)
    }
}

// ---------------------------------------------------------------------------
// Full Model
// ---------------------------------------------------------------------------

/// The complete Qwen3 language model for TTS.
///
/// Architecture:
/// - Token embedding layer
/// - 28 transformer decoder layers
/// - Final RMS normalization
/// - LM head (projects to vocab)
pub struct Qwen3Model {
    embed_tokens: Embedding,
    layers: Vec<Qwen3DecoderLayer>,
    norm: RmsNorm,
    lm_head: Linear,
    rope: RotaryEmbedding,
    config: Qwen3LmConfig,
    /// Last hidden states (before lm_head), used by code predictor.
    last_hidden: std::cell::RefCell<Option<Tensor>>,
}

impl Qwen3Model {
    pub fn new(config: &Qwen3LmConfig, vb: VarBuilder) -> Result<Self> {
        let model_vb = vb.pp("model");

        let embed_tokens = embedding(config.vocab_size, config.hidden_size, model_vb.pp("embed_tokens"))?;

        let mut layers = Vec::with_capacity(config.num_hidden_layers);
        for i in 0..config.num_hidden_layers {
            let layer = Qwen3DecoderLayer::new(config, model_vb.pp(format!("layers.{i}")))?;
            layers.push(layer);
        }

        let norm = rms_norm(config.hidden_size, config.rms_norm_eps, model_vb.pp("norm"))?;

        // LM head — may or may not share weights with embed_tokens
        let lm_head = linear_no_bias(config.hidden_size, config.vocab_size, vb.pp("lm_head"))?;

        let dtype = vb.dtype();
        let device = vb.device().clone();
        let rope = RotaryEmbedding::new(config, dtype, &device)?;

        Ok(Self {
            embed_tokens,
            layers,
            norm,
            lm_head,
            rope,
            config: config.clone(),
            last_hidden: std::cell::RefCell::new(None),
        })
    }

    /// Forward pass through the full model.
    ///
    /// `input_ids`: [batch, seq_len] — token IDs
    /// `kv_caches`: per-layer KV caches
    /// `attention_mask`: optional causal mask [batch, 1, seq_len, total_seq_len]
    ///
    /// Returns logits: [batch, seq_len, vocab_size]
    pub fn forward(
        &self,
        input_ids: &Tensor,
        kv_caches: &mut [KvCache],
        attention_mask: Option<&Tensor>,
    ) -> Result<Tensor> {
        let mut hidden_states = self.embed_tokens.forward(input_ids)?;

        for (i, layer) in self.layers.iter().enumerate() {
            hidden_states =
                layer.forward(&hidden_states, &self.rope, &mut kv_caches[i], attention_mask)?;
        }

        hidden_states = self.norm.forward(&hidden_states)?;

        // Store last hidden state for code predictor
        *self.last_hidden.borrow_mut() = Some(hidden_states.clone());

        let logits = self.lm_head.forward(&hidden_states)?;
        Ok(logits)
    }

    /// Forward pass with pre-computed embeddings (for first iteration where
    /// text embeddings are concatenated with audio features).
    ///
    /// `inputs_embeds`: [batch, seq_len, hidden_size]
    pub fn forward_embeds(
        &self,
        inputs_embeds: &Tensor,
        kv_caches: &mut [KvCache],
        attention_mask: Option<&Tensor>,
    ) -> Result<Tensor> {
        let mut hidden_states = inputs_embeds.clone();

        for (i, layer) in self.layers.iter().enumerate() {
            hidden_states =
                layer.forward(&hidden_states, &self.rope, &mut kv_caches[i], attention_mask)?;
        }

        hidden_states = self.norm.forward(&hidden_states)?;

        *self.last_hidden.borrow_mut() = Some(hidden_states.clone());

        let logits = self.lm_head.forward(&hidden_states)?;
        Ok(logits)
    }

    /// Get the last hidden states (for the code predictor).
    pub fn last_hidden_state(&self) -> Option<Tensor> {
        self.last_hidden.borrow().clone()
    }

    /// Number of transformer layers.
    pub fn num_layers(&self) -> usize {
        self.config.num_hidden_layers
    }

    /// Hidden size.
    pub fn hidden_size(&self) -> usize {
        self.config.hidden_size
    }

    /// Get token embedding layer (for input preparation).
    pub fn embed_tokens(&self) -> &Embedding {
        &self.embed_tokens
    }

    /// Create a causal attention mask.
    pub fn make_causal_mask(
        seq_len: usize,
        past_len: usize,
        dtype: DType,
        device: &Device,
    ) -> Result<Tensor> {
        let total_len = past_len + seq_len;

        if seq_len == 1 {
            // Single token: no masking needed (can attend to everything)
            return Tensor::zeros((1, 1, 1, total_len), dtype, device);
        }

        // Full causal mask: lower triangular
        let mask: Vec<f32> = (0..seq_len)
            .flat_map(|i| {
                (0..total_len).map(move |j| {
                    if j <= past_len + i {
                        0.0
                    } else {
                        f32::NEG_INFINITY
                    }
                })
            })
            .collect();

        Tensor::from_vec(mask, (1, 1, seq_len, total_len), device)?.to_dtype(dtype)
    }
}

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

    #[test]
    fn test_kv_cache() {
        let device = Device::Cpu;
        let mut cache = KvCache::new();
        assert_eq!(cache.seq_len(), 0);

        let k = Tensor::zeros((1, 8, 5, 128), DType::F32, &device).unwrap();
        let v = Tensor::zeros((1, 8, 5, 128), DType::F32, &device).unwrap();
        let (fk, _fv) = cache.append(&k, &v).unwrap();
        assert_eq!(cache.seq_len(), 5);
        assert_eq!(fk.dim(2).unwrap(), 5);

        let k2 = Tensor::zeros((1, 8, 1, 128), DType::F32, &device).unwrap();
        let v2 = Tensor::zeros((1, 8, 1, 128), DType::F32, &device).unwrap();
        let (fk2, _fv2) = cache.append(&k2, &v2).unwrap();
        assert_eq!(cache.seq_len(), 6);
        assert_eq!(fk2.dim(2).unwrap(), 6);

        cache.reset();
        assert_eq!(cache.seq_len(), 0);
    }

    #[test]
    fn test_causal_mask_single_token() {
        let mask = Qwen3Model::make_causal_mask(1, 10, DType::F32, &Device::Cpu).unwrap();
        assert_eq!(mask.dims(), &[1, 1, 1, 11]);
        // All zeros — single token can attend to everything
        let sum: f32 = mask.sum_all().unwrap().to_scalar().unwrap();
        assert_eq!(sum, 0.0);
    }

    #[test]
    fn test_causal_mask_multi_token() {
        let mask = Qwen3Model::make_causal_mask(3, 0, DType::F32, &Device::Cpu).unwrap();
        assert_eq!(mask.dims(), &[1, 1, 3, 3]);
        // Upper triangle should be -inf
        let data: Vec<f32> = mask.flatten_all().unwrap().to_vec1().unwrap();
        // Row 0: [0, -inf, -inf]
        assert_eq!(data[0], 0.0);
        assert!(data[1].is_infinite() && data[1] < 0.0);
        assert!(data[2].is_infinite() && data[2] < 0.0);
        // Row 1: [0, 0, -inf]
        assert_eq!(data[3], 0.0);
        assert_eq!(data[4], 0.0);
        assert!(data[5].is_infinite() && data[5] < 0.0);
        // Row 2: [0, 0, 0]
        assert_eq!(data[6], 0.0);
        assert_eq!(data[7], 0.0);
        assert_eq!(data[8], 0.0);
    }
}