unillm-runtime 0.1.0

//! T5 Model V2 - Clean implementation using solid abstractions
//!
//! This implements the T5 (Text-to-Text Transfer Transformer) architecture including:
//! - T5-small, T5-base, T5-large, T5-3B, T5-11B
//! - Encoder-decoder architecture with relative position embeddings
//! - Bidirectional attention in encoder, causal + cross-attention in decoder

use crate::model_config;
use super::traits::*;
use anyhow::Result;
use serde::{Serialize, Deserialize};

model_config!(T5Config {
    vocab_size: usize = 32128,
    d_model: usize = 512,
    d_kv: usize = 64,
    d_ff: usize = 2048,
    num_layers: usize = 6,
    num_decoder_layers: usize = 6,
    num_heads: usize = 8,
    relative_attention_num_buckets: usize = 32,
    relative_attention_max_distance: usize = 128,
    dropout_rate: f32 = 0.1,
    layer_norm_epsilon: f32 = 1e-6,
    initializer_factor: f32 = 1.0,
    feed_forward_proj: String = "relu".to_string(),
    is_encoder_decoder: bool = true,
    use_cache: bool = true,
    pad_token_id: i64 = 0,
    eos_token_id: i64 = 1,
    decoder_start_token_id: i64 = 0,
    tie_word_embeddings: bool = true,
    is_gated_act: bool = false,
    // Required by model_config macro for ModelConfig trait
    hidden_size: usize = 512,
    num_hidden_layers: usize = 6,
});

impl T5Config {
    /// Create T5Config from GGUF model configuration
    pub fn from_gguf_config(gguf: &crate::weight_loader_core::GGUFModelConfig) -> Self {
        Self {
            vocab_size: gguf.vocab_size,
            d_model: gguf.hidden_size,
            hidden_size: gguf.hidden_size,
            d_kv: gguf.head_dim,
            d_ff: gguf.intermediate_size,
            num_layers: gguf.num_hidden_layers,
            num_decoder_layers: gguf.num_hidden_layers,
            num_hidden_layers: gguf.num_hidden_layers,
            num_heads: gguf.num_attention_heads,
            layer_norm_epsilon: gguf.rms_norm_eps,
            ..Default::default()
        }
    }
}

pub struct T5ModelV2 {
    config: T5Config,
    device: Device,
    shared: Tensor, // Shared embedding table
    encoder: T5Stack,
    decoder: T5Stack,
    lm_head: Option<Tensor>,
}

pub struct T5Stack {
    block: Vec<T5Block>,
    final_layer_norm: Tensor,
    config: T5Config,
    is_decoder: bool,
}

pub struct T5Block {
    self_attention: T5LayerSelfAttention,
    cross_attention: Option<T5LayerCrossAttention>,
    ff: T5LayerFF,
    config: T5Config,
    is_decoder: bool,
}

pub struct T5LayerSelfAttention {
    attention: T5Attention,
    layer_norm: Tensor,
}

pub struct T5LayerCrossAttention {
    attention: T5Attention,
    layer_norm: Tensor,
}

pub struct T5LayerFF {
    wi: Tensor,      // Input projection
    wo: Tensor,      // Output projection
    wi_1: Option<Tensor>, // For gated activations (T5 v1.1)
    layer_norm: Tensor,
    is_gated: bool,
    activation: String,
}

pub struct T5Attention {
    q: Tensor,
    k: Tensor,
    v: Tensor,
    o: Tensor,
    relative_attention_bias: Option<Tensor>,
    num_heads: usize,
    d_kv: usize,
    is_decoder: bool,
    has_relative_attention_bias: bool,
}

impl Model for T5ModelV2 {
    type Config = T5Config;

    fn new(config: T5Config) -> Result<Self> {
        let device = Device::CPU;
        let shared = ops_fn::zeros(&[config.vocab_size, config.d_model], DataType::Float32, &device)?;

        let encoder = T5Stack::new(&config, &device, false)?;
        let decoder = T5Stack::new(&config, &device, true)?;

        let lm_head = if config.tie_word_embeddings {
            None // Will use shared embeddings
        } else {
            Some(ops_fn::zeros(&[config.d_model, config.vocab_size], DataType::Float32, &device)?)
        };

        Ok(Self { config, device, shared, encoder, decoder, lm_head })
    }

    fn from_weights(config: T5Config, weights: ModelWeights) -> Result<Self> {
        let mut model = Self::new(config)?;

        // Load shared embeddings
        if let Some(w) = weights.get("shared.weight") {
            model.shared = w.clone();
        } else if let Some(w) = weights.get("encoder.embed_tokens.weight") {
            model.shared = w.clone();
        }

        // Load lm_head if not tied
        if let Some(w) = weights.get("lm_head.weight") {
            if model.lm_head.is_some() {
                model.lm_head = Some(ops_fn::transpose(w)?);
            }
        }

        // Load encoder and decoder weights
        model.encoder.load_weights(&weights, "encoder")?;
        model.decoder.load_weights(&weights, "decoder")?;

        Ok(model)
    }

    fn forward(&self, inputs: &ModelInputs) -> Result<ModelOutputs> {
        let input_ids = match inputs {
            ModelInputs::Text { input_ids, .. } => input_ids,
            _ => return Err(anyhow::anyhow!("T5 expects text input")),
        };

        // Encoder forward pass
        let encoder_hidden_states = self.encoder.forward(input_ids, &self.shared, None, None)?;

        // For simple forward, use same input for decoder (in practice, decoder input would be shifted)
        let decoder_input = input_ids;

        // Decoder forward pass with cross-attention to encoder outputs
        let decoder_hidden_states = self.decoder.forward(
            decoder_input,
            &self.shared,
            Some(&encoder_hidden_states),
            None,
        )?;

        // Language modeling head: project to vocabulary
        let logits = if let Some(ref lm_head) = self.lm_head {
            ops_fn::matmul(&decoder_hidden_states, lm_head)?
        } else {
            // Tied embeddings: use transposed shared embeddings
            let shared_t = ops_fn::transpose(&self.shared)?;
            ops_fn::matmul(&decoder_hidden_states, &shared_t)?
        };

        Ok(ModelOutputs::Sequence {
            logits,
            encoder_hidden_states: Some(encoder_hidden_states),
            decoder_hidden_states: Some(decoder_hidden_states),
        })
    }

    fn generate(&self, prompt: &str, config: &GenerationConfig) -> Result<String> {
        use crate::tokenizer::Tokenizer;

        // 1. Tokenize input prompt
        let tokenizer = Tokenizer::new();
        let input_tokens: Vec<u32> = tokenizer.encode(prompt);

        // 2. Create encoder input tensor
        let input_i64: Vec<i64> = input_tokens.iter().map(|&t| t as i64).collect();
        let input_tensor = Tensor::from_i64_slice(&input_i64, &[1, input_tokens.len()], &self.device)?;

        // 3. Run encoder once to get hidden states
        let encoder_hidden_states = self.encoder.forward(&input_tensor, &self.shared, None, None)?;

        // 4. Initialize decoder input with decoder_start_token_id
        let mut decoder_tokens: Vec<i64> = vec![self.config.decoder_start_token_id];

        // 5. Autoregressive generation loop
        for _ in 0..config.max_new_tokens {
            // Create decoder input tensor
            let decoder_tensor = Tensor::from_i64_slice(
                &decoder_tokens,
                &[1, decoder_tokens.len()],
                &self.device,
            )?;

            // Run decoder with cached encoder hidden states
            let decoder_hidden_states = self.decoder.forward(
                &decoder_tensor,
                &self.shared,
                Some(&encoder_hidden_states),
                None,
            )?;

            // Get logits for last position
            let logits = if let Some(ref lm_head) = self.lm_head {
                ops_fn::matmul(&decoder_hidden_states, lm_head)?
            } else {
                let shared_t = ops_fn::transpose(&self.shared)?;
                ops_fn::matmul(&decoder_hidden_states, &shared_t)?
            };

            // Extract last token logits and sample
            let logits_candle = logits.to_candle()?;
            let shape = logits_candle.dims();
            let seq_len = if shape.len() == 3 { shape[1] } else { shape[0] };

            let last_logits = if shape.len() == 3 {
                logits_candle.narrow(1, seq_len - 1, 1)?.squeeze(1)?.squeeze(0)?
            } else {
                logits_candle.narrow(0, seq_len - 1, 1)?.squeeze(0)?
            };

            let logits_vec: Vec<f32> = last_logits.to_vec1()?;

            // Greedy sampling (can be extended with temperature, top-k, top-p)
            let next_token = if config.do_sample && config.temperature > 0.0 {
                // Temperature sampling
                let scaled: Vec<f32> = logits_vec.iter()
                    .map(|&x| x / config.temperature)
                    .collect();

                let max_val = scaled.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
                let exp_sum: f32 = scaled.iter().map(|&x| (x - max_val).exp()).sum();
                let probs: Vec<f32> = scaled.iter()
                    .map(|&x| (x - max_val).exp() / exp_sum)
                    .collect();

                use rand::Rng;
                let mut rng = rand::thread_rng();
                let random_val: f32 = rng.gen();
                let mut cumulative = 0.0;
                let mut sampled = 0i64;

                for (idx, &prob) in probs.iter().enumerate() {
                    cumulative += prob;
                    if random_val <= cumulative {
                        sampled = idx as i64;
                        break;
                    }
                }
                sampled
            } else {
                // Greedy
                let mut max_idx = 0;
                let mut max_val = logits_vec[0];
                for (idx, &val) in logits_vec.iter().enumerate() {
                    if val > max_val {
                        max_val = val;
                        max_idx = idx;
                    }
                }
                max_idx as i64
            };

            // Check for EOS
            if next_token == self.config.eos_token_id {
                break;
            }

            decoder_tokens.push(next_token);
        }

        // 6. Decode output tokens (skip decoder_start_token)
        let output_tokens: Vec<u32> = decoder_tokens.iter()
            .skip(1) // Skip decoder_start_token
            .map(|&t| t as u32)
            .collect();

        Ok(tokenizer.decode(&output_tokens))
    }

    fn config(&self) -> &Self::Config {
        &self.config
    }

    fn memory_requirements(&self) -> MemoryRequirements {
        let param_size = (self.config.vocab_size * self.config.d_model +
                         self.config.num_layers * 2 * self.config.d_model * self.config.d_model * 4 +
                         self.config.num_decoder_layers * 2 * self.config.d_model * self.config.d_model * 4) * 4;
        MemoryRequirements {
            gpu_memory: param_size,
            cpu_memory: param_size / 4,
            kv_cache_memory: 2048 * self.config.d_model * 4 * 4, // encoder + decoder KV cache
            peak_memory: param_size + param_size / 2,
        }
    }

    fn to_device(&mut self, device: &Device) -> Result<()> {
        self.device = device.clone();
        self.shared = self.shared.to_device(device)?;
        if let Some(ref mut lm_head) = self.lm_head {
            *lm_head = lm_head.to_device(device)?;
        }
        self.encoder.to_device(device)?;
        self.decoder.to_device(device)?;
        Ok(())
    }
}

impl T5Stack {
    fn new(config: &T5Config, device: &Device, is_decoder: bool) -> Result<Self> {
        let num_layers = if is_decoder {
            config.num_decoder_layers
        } else {
            config.num_layers
        };

        let mut block = Vec::with_capacity(num_layers);
        for i in 0..num_layers {
            // Only first layer has relative position bias
            let has_relative_attention_bias = i == 0;
            block.push(T5Block::new(config, device, is_decoder, has_relative_attention_bias)?);
        }

        let final_layer_norm = ops_fn::zeros(&[config.d_model], DataType::Float32, device)?;

        Ok(Self {
            block,
            final_layer_norm,
            config: config.clone(),
            is_decoder,
        })
    }

    fn forward(
        &self,
        input_ids: &Tensor,
        shared_embedding: &Tensor,
        encoder_hidden_states: Option<&Tensor>,
        attention_mask: Option<&Tensor>,
    ) -> Result<Tensor> {
        // Token embeddings
        let mut hidden_states = ops_fn::embedding(input_ids, shared_embedding)?;

        // Compute position bias from first layer (shared across all layers)
        let shape = hidden_states.shape();
        let seq_len = if shape.len() == 3 { shape[1] } else { shape[0] };
        let position_bias = self.compute_position_bias(seq_len, seq_len)?;

        // Cross-attention position bias (for decoder)
        let cross_position_bias = if self.is_decoder {
            if let Some(enc_hidden) = encoder_hidden_states {
                let enc_shape = enc_hidden.shape();
                let enc_seq_len = if enc_shape.len() == 3 { enc_shape[1] } else { enc_shape[0] };
                Some(self.compute_position_bias(seq_len, enc_seq_len)?)
            } else {
                None
            }
        } else {
            None
        };

        // Apply transformer blocks
        for layer in &self.block {
            hidden_states = layer.forward(
                &hidden_states,
                encoder_hidden_states,
                &position_bias,
                cross_position_bias.as_ref(),
                attention_mask,
            )?;
        }

        // Final layer norm
        ops_fn::layer_norm(&hidden_states, &self.final_layer_norm, None, self.config.layer_norm_epsilon)
    }

    /// Compute relative position bias for T5 attention
    fn compute_position_bias(&self, query_length: usize, key_length: usize) -> Result<Tensor> {
        // For now, return zeros - the actual bias is computed in attention
        // This is a placeholder for the relative position encoding
        ops_fn::zeros(
            &[1, self.config.num_heads, query_length, key_length],
            DataType::Float32,
            &Device::CPU,
        )
    }

    fn load_weights(&mut self, weights: &ModelWeights, prefix: &str) -> Result<()> {
        // Load final layer norm
        if let Some(w) = weights.get(&format!("{}.final_layer_norm.weight", prefix)) {
            self.final_layer_norm = w.clone();
        }

        // Load block weights
        for (i, block) in self.block.iter_mut().enumerate() {
            block.load_weights(weights, &format!("{}.block.{}", prefix, i))?;
        }

        Ok(())
    }

    fn to_device(&mut self, device: &Device) -> Result<()> {
        self.final_layer_norm = self.final_layer_norm.to_device(device)?;
        for block in &mut self.block {
            block.to_device(device)?;
        }
        Ok(())
    }
}

impl T5Block {
    fn new(config: &T5Config, device: &Device, is_decoder: bool, has_relative_attention_bias: bool) -> Result<Self> {
        // Self-attention layer
        let self_attention = T5LayerSelfAttention::new(config, device, is_decoder, has_relative_attention_bias)?;

        // Cross-attention layer (decoder only)
        let cross_attention = if is_decoder {
            Some(T5LayerCrossAttention::new(config, device, has_relative_attention_bias)?)
        } else {
            None
        };

        // Feed-forward layer
        let ff = T5LayerFF::new(config, device)?;

        Ok(Self {
            self_attention,
            cross_attention,
            ff,
            config: config.clone(),
            is_decoder,
        })
    }

    fn forward(
        &self,
        hidden_states: &Tensor,
        encoder_hidden_states: Option<&Tensor>,
        position_bias: &Tensor,
        cross_position_bias: Option<&Tensor>,
        attention_mask: Option<&Tensor>,
    ) -> Result<Tensor> {
        // 1. Self-attention with residual
        let attn_output = self.self_attention.forward(
            hidden_states,
            hidden_states,
            position_bias,
            attention_mask,
        )?;
        let hidden_states = ops_fn::add(hidden_states, &attn_output)?;

        // 2. Cross-attention (decoder only) with residual
        let hidden_states = if let (Some(cross_attn), Some(enc_hidden)) = (&self.cross_attention, encoder_hidden_states) {
            let cross_output = cross_attn.forward(
                &hidden_states,
                enc_hidden,
                cross_position_bias.unwrap_or(position_bias),
                None, // No mask for cross-attention typically
            )?;
            ops_fn::add(&hidden_states, &cross_output)?
        } else {
            hidden_states
        };

        // 3. Feed-forward with residual
        let ff_output = self.ff.forward(&hidden_states)?;
        ops_fn::add(&hidden_states, &ff_output)
    }

    fn load_weights(&mut self, weights: &ModelWeights, prefix: &str) -> Result<()> {
        // Load self-attention weights
        self.self_attention.load_weights(weights, &format!("{}.layer.0", prefix))?;

        // Load cross-attention weights (decoder only)
        if let Some(ref mut cross_attn) = self.cross_attention {
            cross_attn.load_weights(weights, &format!("{}.layer.1", prefix))?;
        }

        // Load feed-forward weights
        let ff_layer_idx = if self.is_decoder { 2 } else { 1 };
        self.ff.load_weights(weights, &format!("{}.layer.{}", prefix, ff_layer_idx))?;

        Ok(())
    }

    fn to_device(&mut self, device: &Device) -> Result<()> {
        self.self_attention.to_device(device)?;
        if let Some(ref mut cross_attn) = self.cross_attention {
            cross_attn.to_device(device)?;
        }
        self.ff.to_device(device)?;
        Ok(())
    }
}

impl T5LayerSelfAttention {
    fn new(config: &T5Config, device: &Device, is_decoder: bool, has_relative_attention_bias: bool) -> Result<Self> {
        let attention = T5Attention::new(config, device, is_decoder, has_relative_attention_bias)?;
        let layer_norm = ops_fn::zeros(&[config.d_model], DataType::Float32, device)?;

        Ok(Self { attention, layer_norm })
    }

    fn forward(
        &self,
        hidden_states: &Tensor,
        key_value_states: &Tensor,
        position_bias: &Tensor,
        attention_mask: Option<&Tensor>,
    ) -> Result<Tensor> {
        // Pre-LayerNorm
        let normed = ops_fn::layer_norm(hidden_states, &self.layer_norm, None, 1e-6)?;

        // Self-attention
        self.attention.forward(&normed, key_value_states, position_bias, attention_mask)
    }

    fn load_weights(&mut self, weights: &ModelWeights, prefix: &str) -> Result<()> {
        // Layer norm
        if let Some(w) = weights.get(&format!("{}.layer_norm.weight", prefix)) {
            self.layer_norm = w.clone();
        }

        // Attention weights
        self.attention.load_weights(weights, &format!("{}.SelfAttention", prefix))?;

        Ok(())
    }

    fn to_device(&mut self, device: &Device) -> Result<()> {
        self.layer_norm = self.layer_norm.to_device(device)?;
        self.attention.to_device(device)?;
        Ok(())
    }
}

impl T5LayerCrossAttention {
    fn new(config: &T5Config, device: &Device, has_relative_attention_bias: bool) -> Result<Self> {
        // Cross-attention is never causal
        let attention = T5Attention::new(config, device, false, has_relative_attention_bias)?;
        let layer_norm = ops_fn::zeros(&[config.d_model], DataType::Float32, device)?;

        Ok(Self { attention, layer_norm })
    }

    fn forward(
        &self,
        hidden_states: &Tensor,
        key_value_states: &Tensor,
        position_bias: &Tensor,
        attention_mask: Option<&Tensor>,
    ) -> Result<Tensor> {
        // Pre-LayerNorm
        let normed = ops_fn::layer_norm(hidden_states, &self.layer_norm, None, 1e-6)?;

        // Cross-attention (query from decoder, key/value from encoder)
        self.attention.forward(&normed, key_value_states, position_bias, attention_mask)
    }

    fn load_weights(&mut self, weights: &ModelWeights, prefix: &str) -> Result<()> {
        // Layer norm
        if let Some(w) = weights.get(&format!("{}.layer_norm.weight", prefix)) {
            self.layer_norm = w.clone();
        }

        // Attention weights
        self.attention.load_weights(weights, &format!("{}.EncDecAttention", prefix))?;

        Ok(())
    }

    fn to_device(&mut self, device: &Device) -> Result<()> {
        self.layer_norm = self.layer_norm.to_device(device)?;
        self.attention.to_device(device)?;
        Ok(())
    }
}

impl T5LayerFF {
    fn new(config: &T5Config, device: &Device) -> Result<Self> {
        let is_gated = config.is_gated_act ||
                       config.feed_forward_proj.contains("gated");

        let wi = ops_fn::zeros(&[config.d_model, config.d_ff], DataType::Float32, device)?;
        let wo = ops_fn::zeros(&[config.d_ff, config.d_model], DataType::Float32, device)?;

        let wi_1 = if is_gated {
            Some(ops_fn::zeros(&[config.d_model, config.d_ff], DataType::Float32, device)?)
        } else {
            None
        };

        let layer_norm = ops_fn::zeros(&[config.d_model], DataType::Float32, device)?;

        // Determine activation function
        let activation = if config.feed_forward_proj.contains("gelu") {
            "gelu".to_string()
        } else {
            "relu".to_string()
        };

        Ok(Self {
            wi,
            wo,
            wi_1,
            layer_norm,
            is_gated,
            activation,
        })
    }

    fn forward(&self, hidden_states: &Tensor) -> Result<Tensor> {
        // Pre-LayerNorm
        let normed = ops_fn::layer_norm(hidden_states, &self.layer_norm, None, 1e-6)?;

        // Feed-forward computation
        let hidden = if self.is_gated {
            // Gated activation: gate * activation(up)
            let gate = ops_fn::matmul(&normed, &self.wi)?;
            let up = ops_fn::matmul(&normed, self.wi_1.as_ref().unwrap())?;

            let activated = match self.activation.as_str() {
                "gelu" => ops_fn::gelu(&gate)?,
                _ => relu(&gate)?,
            };

            ops_fn::mul(&activated, &up)?
        } else {
            // Standard: activation(up)
            let up = ops_fn::matmul(&normed, &self.wi)?;
            match self.activation.as_str() {
                "gelu" => ops_fn::gelu(&up)?,
                _ => relu(&up)?,
            }
        };

        // Down projection
        ops_fn::matmul(&hidden, &self.wo)
    }

    fn load_weights(&mut self, weights: &ModelWeights, prefix: &str) -> Result<()> {
        // Layer norm
        if let Some(w) = weights.get(&format!("{}.layer_norm.weight", prefix)) {
            self.layer_norm = w.clone();
        }

        // Dense relu dense weights (transpose for matmul)
        if let Some(w) = weights.get(&format!("{}.DenseReluDense.wi.weight", prefix)) {
            self.wi = ops_fn::transpose(w)?;
        }
        // Alternative naming for gated variants
        if let Some(w) = weights.get(&format!("{}.DenseReluDense.wi_0.weight", prefix)) {
            self.wi = ops_fn::transpose(w)?;
        }
        if let Some(w) = weights.get(&format!("{}.DenseReluDense.wi_1.weight", prefix)) {
            self.wi_1 = Some(ops_fn::transpose(w)?);
        }
        if let Some(w) = weights.get(&format!("{}.DenseReluDense.wo.weight", prefix)) {
            self.wo = ops_fn::transpose(w)?;
        }

        Ok(())
    }

    fn to_device(&mut self, device: &Device) -> Result<()> {
        self.wi = self.wi.to_device(device)?;
        self.wo = self.wo.to_device(device)?;
        if let Some(ref mut wi_1) = self.wi_1 {
            *wi_1 = wi_1.to_device(device)?;
        }
        self.layer_norm = self.layer_norm.to_device(device)?;
        Ok(())
    }
}

impl T5Attention {
    fn new(config: &T5Config, device: &Device, is_decoder: bool, has_relative_attention_bias: bool) -> Result<Self> {
        let inner_dim = config.num_heads * config.d_kv;

        // Q, K, V, O projections
        let q = ops_fn::zeros(&[config.d_model, inner_dim], DataType::Float32, device)?;
        let k = ops_fn::zeros(&[config.d_model, inner_dim], DataType::Float32, device)?;
        let v = ops_fn::zeros(&[config.d_model, inner_dim], DataType::Float32, device)?;
        let o = ops_fn::zeros(&[inner_dim, config.d_model], DataType::Float32, device)?;

        // Relative position bias (only first layer of each stack)
        let relative_attention_bias = if has_relative_attention_bias {
            Some(ops_fn::zeros(
                &[config.relative_attention_num_buckets, config.num_heads],
                DataType::Float32,
                device,
            )?)
        } else {
            None
        };

        Ok(Self {
            q,
            k,
            v,
            o,
            relative_attention_bias,
            num_heads: config.num_heads,
            d_kv: config.d_kv,
            is_decoder,
            has_relative_attention_bias,
        })
    }

    fn forward(
        &self,
        hidden_states: &Tensor,
        key_value_states: &Tensor,
        position_bias: &Tensor,
        attention_mask: Option<&Tensor>,
    ) -> Result<Tensor> {
        let shape = hidden_states.shape();
        let (batch_size, seq_len) = if shape.len() == 3 {
            (shape[0], shape[1])
        } else {
            (1, shape[0])
        };

        let kv_shape = key_value_states.shape();
        let kv_seq_len = if kv_shape.len() == 3 { kv_shape[1] } else { kv_shape[0] };

        // Project to Q, K, V
        let query = ops_fn::matmul(hidden_states, &self.q)?;
        let key = ops_fn::matmul(key_value_states, &self.k)?;
        let value = ops_fn::matmul(key_value_states, &self.v)?;

        // Reshape for multi-head attention: [batch, seq, heads * d_kv] -> [batch, heads, seq, d_kv]
        let q_candle = query.to_candle()?;
        let k_candle = key.to_candle()?;
        let v_candle = value.to_candle()?;

        let q_reshaped = q_candle
            .reshape(&[batch_size, seq_len, self.num_heads, self.d_kv])?
            .transpose(1, 2)?;

        let k_reshaped = k_candle
            .reshape(&[batch_size, kv_seq_len, self.num_heads, self.d_kv])?
            .transpose(1, 2)?;

        let v_reshaped = v_candle
            .reshape(&[batch_size, kv_seq_len, self.num_heads, self.d_kv])?
            .transpose(1, 2)?;

        // Compute attention scores: Q @ K^T
        let k_t = k_reshaped.transpose(2, 3)?;
        let scores = q_reshaped.contiguous()?.matmul(&k_t.contiguous()?)?;

        // Add relative position bias
        let position_bias_candle = position_bias.to_candle()?;
        let scores = scores.broadcast_add(&position_bias_candle)?;

        // Apply causal mask for decoder self-attention
        let scores = if self.is_decoder && seq_len == kv_seq_len {
            // Create causal mask
            let device = scores.device();
            let mut mask_data = vec![0.0f32; seq_len * kv_seq_len];
            for i in 0..seq_len {
                for j in 0..kv_seq_len {
                    if j > i {
                        mask_data[i * kv_seq_len + j] = f32::NEG_INFINITY;
                    }
                }
            }
            let causal_mask = candle_core::Tensor::from_vec(mask_data, &[1, 1, seq_len, kv_seq_len], device)?;
            scores.broadcast_add(&causal_mask)?
        } else {
            scores
        };

        // Apply attention mask if provided
        let scores = if let Some(mask) = attention_mask {
            let mask_candle = mask.to_candle()?;
            scores.broadcast_add(&mask_candle)?
        } else {
            scores
        };

        // Softmax
        let attention_weights = candle_nn::ops::softmax_last_dim(&scores)?;

        // Apply attention to values
        let attn_output = attention_weights.matmul(&v_reshaped.contiguous()?)?;

        // Reshape back: [batch, heads, seq, d_kv] -> [batch, seq, heads * d_kv]
        let attn_output = attn_output
            .transpose(1, 2)?
            .reshape(&[batch_size, seq_len, self.num_heads * self.d_kv])?;

        let attn_output = Tensor::from_candle(attn_output);

        // Output projection
        ops_fn::matmul(&attn_output, &self.o)
    }

    fn load_weights(&mut self, weights: &ModelWeights, prefix: &str) -> Result<()> {
        // Load and transpose projection weights
        if let Some(w) = weights.get(&format!("{}.q.weight", prefix)) {
            self.q = ops_fn::transpose(w)?;
        }
        if let Some(w) = weights.get(&format!("{}.k.weight", prefix)) {
            self.k = ops_fn::transpose(w)?;
        }
        if let Some(w) = weights.get(&format!("{}.v.weight", prefix)) {
            self.v = ops_fn::transpose(w)?;
        }
        if let Some(w) = weights.get(&format!("{}.o.weight", prefix)) {
            self.o = ops_fn::transpose(w)?;
        }

        // Load relative attention bias
        if let Some(ref mut bias) = self.relative_attention_bias {
            if let Some(w) = weights.get(&format!("{}.relative_attention_bias.weight", prefix)) {
                *bias = w.clone();
            }
        }

        Ok(())
    }

    fn to_device(&mut self, device: &Device) -> Result<()> {
        self.q = self.q.to_device(device)?;
        self.k = self.k.to_device(device)?;
        self.v = self.v.to_device(device)?;
        self.o = self.o.to_device(device)?;

        if let Some(ref mut bias) = self.relative_attention_bias {
            *bias = bias.to_device(device)?;
        }

        Ok(())
    }
}

/// ReLU activation function
fn relu(input: &Tensor) -> Result<Tensor> {
    let x = input.to_candle()?;
    let result = x.relu()?;
    Ok(Tensor::from_candle(result))
}

/// Compute relative position bucket for T5 attention
/// This converts relative positions to bucket indices for the position bias lookup
#[allow(dead_code)]
fn relative_position_bucket(
    relative_position: i32,
    bidirectional: bool,
    num_buckets: usize,
    max_distance: usize,
) -> usize {
    let mut relative_buckets = 0usize;
    let mut relative_position = relative_position;

    if bidirectional {
        let num_buckets = num_buckets / 2;
        if relative_position > 0 {
            relative_buckets = num_buckets;
        } else {
            relative_position = -relative_position;
        }
    } else {
        relative_position = (-relative_position).max(0);
    }

    let relative_position = relative_position as usize;

    // Half of buckets are for exact positions
    let max_exact = num_buckets / 2;

    if relative_position < max_exact {
        relative_buckets + relative_position
    } else {
        // The other half are for logarithmically bigger bins
        let relative_position_if_large = max_exact +
            ((relative_position as f32 / max_exact as f32).ln() /
             (max_distance as f32 / max_exact as f32).ln() *
             (num_buckets - max_exact) as f32) as usize;
        relative_buckets + relative_position_if_large.min(num_buckets - 1)
    }
}

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

    #[test]
    fn test_t5_model_creation() {
        let config = T5Config {
            vocab_size: 1000,
            d_model: 64,
            hidden_size: 64,
            d_kv: 16,
            d_ff: 256,
            num_layers: 2,
            num_decoder_layers: 2,
            num_hidden_layers: 2,
            num_heads: 4,
            ..Default::default()
        };

        let model = T5ModelV2::new(config).unwrap();
        assert_eq!(model.config().vocab_size(), 1000);
        assert_eq!(model.config().hidden_size(), 64);
        assert_eq!(model.config().num_layers(), 2);
    }

    #[test]
    fn test_t5_forward_pass() {
        let config = T5Config {
            vocab_size: 100,
            d_model: 32,
            hidden_size: 32,
            d_kv: 8,
            d_ff: 64,
            num_layers: 1,
            num_decoder_layers: 1,
            num_hidden_layers: 1,
            num_heads: 4,
            ..Default::default()
        };

        let model = T5ModelV2::new(config).unwrap();
        let input_ids = ops_fn::zeros(&[2, 8], DataType::Int64, &Device::CPU).unwrap();
        let inputs = ModelInputs::text(input_ids);

        let outputs = model.forward(&inputs).unwrap();
        match outputs {
            ModelOutputs::Sequence { logits, encoder_hidden_states, decoder_hidden_states } => {
                assert_eq!(logits.shape(), &[2, 8, 100]); // batch, seq, vocab
                assert!(encoder_hidden_states.is_some());
                assert!(decoder_hidden_states.is_some());
            }
            _ => panic!("Expected sequence output"),
        }
    }

    #[test]
    fn test_t5_generation() {
        let config = T5Config {
            vocab_size: 256,
            d_model: 32,
            hidden_size: 32,
            d_kv: 8,
            d_ff: 64,
            num_layers: 1,
            num_decoder_layers: 1,
            num_hidden_layers: 1,
            num_heads: 4,
            ..Default::default()
        };
        let model = T5ModelV2::new(config).unwrap();
        let gen_config = GenerationConfig {
            max_new_tokens: 5,
            ..Default::default()
        };

        let output = model.generate("Hello", &gen_config).unwrap();
        // Should produce some output (even if random with uninitialized weights)
        // Generation may produce empty output if EOS token is sampled early
        let _ = output;
    }

    #[test]
    fn test_relative_position_bucket() {
        // Test bidirectional bucketing (encoder)
        let bucket = relative_position_bucket(0, true, 32, 128);
        assert_eq!(bucket, 0);

        let bucket = relative_position_bucket(1, true, 32, 128);
        assert!(bucket > 0);

        let bucket = relative_position_bucket(-1, true, 32, 128);
        assert!(bucket < 16); // Should be in first half

        // Test unidirectional bucketing (decoder)
        let bucket = relative_position_bucket(0, false, 32, 128);
        assert_eq!(bucket, 0);
    }
}