Blame for src/atencao.rs

//! ADBA — Atenção Dual-Bloco Adaptativa (APRIMORAMENTO 1 do usuário:
//! "aprimoramento de Attention Mechanism durante o processamento paralelo
//! de NCM"). Prova completa na memória r-04; síntese:
//!
//! (i)  [paralelismo preservado] o resumo-âncora por nível h_ℓ é função
//!      APENAS de c_{<ℓ} (computado do cache/entrada corrente), logo a
//!      fatoração p(c_ℓ|c_{<ℓ}) = Π_i p(c_{ℓ,i}|c_{<ℓ}) e NFE = L
//!      permanecem EXATOS (Teo B.2/D.1);
//! (ii) [expressividade ⊇ NCM] com γ=0, w_ℓ'=1, τ_ℓ=1 a ADBA reduz-se à
//!      atenção block-causal do NCM (classe estritamente contida);
//! (iii)[planejamento reforçado] o âncora agrega o estado global dos níveis
//!      anteriores num vetor por nível — injetado por FiLM — fortalecendo a
//!      hipótese H2 do Teo C.2 (condicionais consistentes dado o plano).
//!
//! Componentes: (a) atenção multi-cabeça com máscara aditiva arbitrária +
//! KV-cache; (b) reponderação aprendível w_{ℓ'} dos logits por nível-fonte;
//! (c) temperatura aprendível τ_ℓ por nível; (d) injeção FiLM do resumo.

use candle_core::{Result, Tensor, Var, D};
use crate::camadas::{Linear, RmsNorm, Swiglu};

pub struct CabecaAdba {
    pub q: Linear,
    pub k: Linear,
    pub v: Linear,
    pub o: Linear,
    /// Reponderação de logits por nível-fonte (len = nº de níveis + 1 p/ prefixo).
    pub w_nivel: Option<Var>,
    /// Temperatura aprendível por nível (log-espaço, neutro = 0).
    pub log_tau_nivel: Option<Var>,
    pub n_cabecas: usize,
    pub d_cabeca: usize,
}

#[derive(Default, Clone)]
pub struct CacheKv {
    pub k: Option<Tensor>,
    pub v: Option<Tensor>,
}

impl CabecaAdba {
    pub fn nova(d_modelo: usize, n_cabecas: usize, n_niveis: usize, adba_ativo: bool, reponderar: bool, tau_nivel: bool, dev: &candle_core::Device) -> Result<Self> {
        let d_cabeca = d_modelo / n_cabecas;
        let desvio = 1.0 / ((d_modelo * 2) as f64).sqrt();
        let mk = || Linear::sem_vies(d_modelo, d_modelo, desvio, dev);
        let w_nivel = if adba_ativo && reponderar {
            Some(Var::from_tensor(&Tensor::ones((n_niveis + 1,), candle_core::DType::F32, dev)?)?)
        } else {
            None
        };
        let log_tau_nivel = if adba_ativo && tau_nivel {
            Some(Var::from_tensor(&Tensor::zeros((n_niveis + 1,), candle_core::DType::F32, dev)?)?)
        } else {
            None
        };
        Ok(CabecaAdba {
            q: mk()?, k: mk()?, v: mk()?, o: mk()?,
            w_nivel, log_tau_nivel,
            n_cabecas, d_cabeca,
        })
    }

    fn dividir_cabecas(&self, x: &Tensor) -> Result<Tensor> {
        let (b, n, _) = (x.dim(0)?, x.dim(1)?, x.dim(2)?);
        x.reshape((b, n, self.n_cabecas, self.d_cabeca))?
            .transpose(1, 2)?
            .contiguous()
    }

    fn juntar_cabecas(&self, x: &Tensor) -> Result<Tensor> {
        let b = x.dim(0)?;
        let n = x.dim(2)?;
        x.transpose(1, 2)?.contiguous()?
            .reshape((b, n, self.n_cabecas * self.d_cabeca))
    }

    /// Atenção com máscara aditiva (T_k × T_k) e offsets de nível para a
    /// reponderação/temperatura da ADBA. `niveis_query`/`niveis_chave` dizem
    /// a qual nível (0=prefixo, 1..L) pertence cada posição.
    pub fn attention(
        &self,
        x: &Tensor,
        mascara: Option<&Tensor>,
        cache: &mut CacheKv,
        niveis_query: Option<&[usize]>,
        niveis_chave: Option<&[usize]>,
    ) -> Result<Tensor> {
        let b = x.dim(0)?;
        let q = self.dividir_cabecas(&self.q.fwd(x)?)?;
        let k = self.dividir_cabecas(&self.k.fwd(x)?)?;
        let v = self.dividir_cabecas(&self.v.fwd(x)?)?;
        let (k, v) = match (&cache.k, &cache.v) {
            (Some(pk), Some(pv)) => (Tensor::cat(&[pk, &k], 2)?, Tensor::cat(&[pv, &v], 2)?),
            _ => (k, v),
        };
        cache.k = Some(k.clone());
        cache.v = Some(v.clone());
        let escala = 1.0 / (self.d_cabeca as f64).sqrt();
        let mut scores = q.matmul(&k.transpose(2, 3)?)?.affine(escala, 0.0)?;
        // ADBA (ii): reponderação por nível-fonte — score += log w_{ℓ'}
        if let (Some(w), Some(nq), Some(nk)) = (&self.w_nivel, niveis_query, niveis_chave) {
            let t_k = scores.dim(3)?;
            let w_clamped = w.as_tensor().clamp(1e-3f32, 1e3f32)?;
            let log_w = (&w_clamped + 1e-8)?.log()?;
            let mut biases = vec![0f32; t_k];
            for (j, &nivel) in nk.iter().enumerate().take(t_k) {
                biases[j] = log_w.get(nivel)?.to_scalar::<f32>()?;
            }
            let _ = nq;
            let bias = Tensor::from_vec(biases, (t_k,), x.device())?;
            scores = scores.broadcast_add(&bias.reshape((1, 1, 1, t_k))?)?;
        }
        // ADBA (iii): temperatura por nível do QUERY
        if let (Some(lt), Some(nq), _) = (&self.log_tau_nivel, niveis_query, niveis_chave) {
            let t_q = scores.dim(2)?;
            let tau = lt.as_tensor().affine(0.1, 0.0)?.exp()?; // suave (±~10% por unidade)
            let mut fatores = vec![1f32; t_q];
            for (i, &nivel) in nq.iter().enumerate().take(t_q) {
                fatores[i] = tau.get(nivel)?.to_scalar::<f32>()?;
            }
            let f = Tensor::from_vec(fatores, (t_q,), x.device())?
                .reshape((1, 1, t_q, 1))?;
            scores = scores.broadcast_div(&f)?;
        }
        if let Some(m) = mascara {
            scores = scores.broadcast_add(m)?;
        }
        let attn = candle_nn::ops::softmax(&scores, D::Minus1)?;
        let ctx = attn.matmul(&v)?;
        self.juntar_cabecas(&ctx)
    }

    pub fn vars(&self) -> Vec<Var> {
        let mut v: Vec<Var> = [self.q.vars(), self.k.vars(), self.v.vars(), self.o.vars()].concat();
        if let Some(w) = &self.w_nivel { v.push(w.clone()); }
        if let Some(t) = &self.log_tau_nivel { v.push(t.clone()); }
        v
    }
}

/// Bloco ADBA completo: pré-norm + atenção + resíduo + pré-norm + SwiGLU.
pub struct BlocoAdba {
    pub norm1: RmsNorm,
    pub attn: CabecaAdba,
    pub norm2: RmsNorm,
    pub ffn: Swiglu,
}

impl BlocoAdba {
    pub fn novo(d_modelo: usize, n_cabecas: usize, d_ff: usize, n_niveis: usize, adba_ativo: bool, reponderar: bool, tau_nivel: bool, dev: &candle_core::Device) -> Result<Self> {
        Ok(BlocoAdba {
            norm1: RmsNorm::novo(d_modelo, 1e-5, dev)?,
            attn: CabecaAdba::nova(d_modelo, n_cabecas, n_niveis, adba_ativo, reponderar, tau_nivel, dev)?,
            norm2: RmsNorm::novo(d_modelo, 1e-5, dev)?,
            ffn: Swiglu::novo(d_modelo, d_ff, dev)?,
        })
    }

    pub fn fwd(
        &self,
        x: &Tensor,
        mascara: Option<&Tensor>,
        cache: &mut CacheKv,
        niveis_query: Option<&[usize]>,
        niveis_chave: Option<&[usize]>,
    ) -> Result<Tensor> {
        let h = self.attn.attention(&self.norm1.fwd(x)?, mascara, cache, niveis_query, niveis_chave)?;
        let h = self.attn.o.fwd(&h)?;
        let x = x.add(&h)?;
        let y = self.ffn.fwd(&self.norm2.fwd(&x)?)?;
        x.add(&y)
    }

    pub fn vars(&self) -> Vec<Var> {
        [self.norm1.vars(), self.attn.vars(), self.norm2.vars(), self.ffn.vars()].concat()
    }
}

/// Resumo-âncora por nível (ADBA i): h_ℓ = Attn(q_ℓ, K/V = c_{<ℓ}) — pooling
/// sobre as posições dos níveis anteriores (função de c_{<ℓ} APENAS).
pub struct AncoraNivel {
    pub q: Var, // (d,) vetor de consulta aprendível por nível — (L, d)
    pub proj: Linear,
}

impl AncoraNivel {
    pub fn nova(d: usize, n_niveis: usize, dev: &candle_core::Device) -> Result<Self> {
        let q = Tensor::randn(0f32, 0.02, (n_niveis, d), dev)?;
        Ok(AncoraNivel { q: Var::from_tensor(&q)?, proj: Linear::novo(d, d, 0.02, dev)? })
    }

    /// contexto (1, n_ctx, d) das posições dos níveis < ℓ → resumo (1, d).
    pub fn resumir(&self, contexto: &Tensor, nivel: usize) -> Result<Tensor> {
        let d = self.q.as_tensor().dim(1)?;
        let q = self.q.as_tensor().get(nivel)?.reshape((1, 1, d))?;
        // atenção de cabeça única: softmax(q·Kᵀ/√d)·V
        let escala = 1.0 / (d as f64).sqrt();
        let scores = q.matmul(&contexto.transpose(1, 2)?)?.affine(escala, 0.0)?;
        let pesos = candle_nn::ops::softmax(&scores, D::Minus1)?;
        let resumo = pesos.matmul(contexto)?.reshape((1, d))?;
        self.proj.fwd(&resumo)
    }

    pub fn vars(&self) -> Vec<Var> {
        let mut v = vec![self.q.clone()];
        v.extend(self.proj.vars());
        v
    }
}