Blame for src/amostrador.rs

//! Amostrador NCM (inferência) — port de `conceitos/amostrador.py` (n-04 §2.5).
//!
//! Laço de níveis (Teo B.2/B.4): EXATAMENTE L avaliações do gerador
//! (NFE = L — V10; com CFG ativa são 2L — Cor. B.2.1, ainda ≪ T·N da difusão);
//! parada anytime em k (Teo A.2/B.4); temperatura-auto por nível por bisseção
//! de entropia (Teo D.3/28.8 — monotonia estrita de H em τ); CFG por
//! interpolação de logits entre ramos condicionado/não-condicionado (Teo
//! D.2(iii)); MBR com utilidade de concordância (Teo C.4/T6).

use candle_core::{Device, Result, Tensor, D};
use crate::gerador::{GeradorConceitos, N_PREFIXO};
use crate::tokenizador::TokenizadorConceitos;
use crate::tronco::CacheKv;
use rand::prelude::*;

pub struct AmostradorConceitos {
    pub gerador: *mut GeradorConceitos,
    pub tok: *const TokenizadorConceitos,
    pub nfe_contador: usize,
    pub temperaturas: Vec<f64>,
    pub limites_tau: [f64; 2],
    pub cfg_scale: f64,
    pub entropia_alvo: Option<f64>,
}

unsafe impl Send for AmostradorConceitos {}

impl AmostradorConceitos {
    pub fn novo(gerador: &mut GeradorConceitos, tok: &TokenizadorConceitos, cfg_scale: f64, limites: [f64; 2], entropia_alvo: Option<f64>) -> Self {
        let l = gerador.niveis.len();
        AmostradorConceitos {
            gerador: gerador as *mut _,
            tok: tok as *const _,
            nfe_contador: 0,
            temperaturas: vec![1.0; l],
            limites_tau: limites,
            cfg_scale,
            entropia_alvo,
        }
    }

    unsafe fn g(&mut self) -> &'static mut GeradorConceitos {
        &mut *self.gerador
    }
    unsafe fn t(&self) -> &'static TokenizadorConceitos {
        &*self.tok
    }

    /// Bisseção de entropia-alvo para τ_ℓ SOBRE OS LOGITS (Teo D.3(ii)):
    /// H(softmax(logits/τ)) é estritamente crescente em τ (lema do Teo 28.8).
    pub fn calibrar_tau(&self, logits: &Tensor, iters: usize) -> Result<f64> {
        let alvo = match self.entropia_alvo {
            Some(h) => h,
            None => return Ok(1.0),
        };
        let (mut lo, mut hi) = (self.limites_tau[0], self.limites_tau[1]);
        for _ in 0..iters {
            let meio = 0.5 * (lo + hi);
            let h = entropia_normalizada_de_logits(logits, meio)?;
            if h < alvo { lo = meio; } else { hi = meio; }
        }
        Ok(0.5 * (lo + hi))
    }

    /// Gera UMA hierarquia. NFE = L (ou 2L com CFG — contado exatamente).
    /// `parar_no_nivel`: anytime (Teo B.4). Devolve (códigos por nível, NFE).
    pub fn gerar(
        &mut self,
        b: usize,
        dev: &Device,
        rng: &mut StdRng,
        parar_no_nivel: Option<usize>,
        greedy: bool,
    ) -> Result<(Vec<Tensor>, usize)> {
        unsafe {
            let g = &mut *self.gerador;
            let t = &*self.tok;
            let l = g.niveis.len();
            let cfg_ativo = (self.cfg_scale - 1.0).abs() > 1e-9;
            let mut caches_cond: Vec<CacheKv> = Vec::new();
            let mut caches_uncond: Vec<CacheKv> = Vec::new();
            let mut zhats: Vec<Tensor> = Vec::new();
            let mut codigos: Vec<Tensor> = Vec::new();
            // v070-2 (paridade FiLM): estados pós-bloco-0 acumulados por ramo —
            // o âncora mede o resumo sobre eles exatamente como forward_treino
            // (que usa o x PRÉ-FiLM pós-bloco-0 — vide cat tardio em tronco.rs).
            let d_modelo = g.tronco.norm_final.ganho.as_tensor().dims()[0];
            let mut estados_b0_cond = Tensor::zeros((b, 0, d_modelo), candle_core::DType::F32, dev)?;
            let mut estados_b0_uncond = Tensor::zeros((b, 0, d_modelo), candle_core::DType::F32, dev)?;
            let fim = parar_no_nivel.unwrap_or(l).min(l);
            let nfe_inicio = self.nfe_contador;
            for nivel in 0..fim {
                let n_l = g.niveis[nivel];
                let entrada_cond = g.entrada_nivel_pub(zhats.last(), nivel, b, dev)?;
                // ramo não-condicionado: mesmo nível com ẑ = 0 (D.2(iii))
                let entrada_uncond = if cfg_ativo {
                    Some(g.entrada_nivel_pub(None, nivel, b, dev)?)
                } else {
                    None
                };
                // resumo-âncora para o FiLM deste nível (nível ≥ 2) — paridade
                // com forward_treino: q[nivel] sobre estados_b0 acumulados;
                // nível 1: resumo ≡ 0 ⇒ FiLM identidade (W_f(0)=0) ⇒ None.
                let film_cond: Option<(Tensor, usize)> = if nivel >= 1 {
                    match &g.tronco.ancora {
                        Some(anc) => {
                            let mut r_nivel = Vec::new();
                            for bi in 0..b {
                                r_nivel.push(anc.resumir(
                                    &estados_b0_cond.narrow(0, bi, 1)?, nivel)?);
                            }
                            Some((Tensor::cat(&r_nivel, 0)?, nivel + 1))
                        }
                        None => None,
                    }
                } else { None };
                let film_uncond: Option<(Tensor, usize)> = if nivel >= 1 && cfg_ativo {
                    match &g.tronco.ancora {
                        Some(anc) => {
                            let mut r_nivel = Vec::new();
                            for bi in 0..b {
                                r_nivel.push(anc.resumir(
                                    &estados_b0_uncond.narrow(0, bi, 1)?, nivel)?);
                            }
                            Some((Tensor::cat(&r_nivel, 0)?, nivel + 1))
                        }
                        None => None,
                    }
                } else { None };
                let (h_c, h_u): (Tensor, Option<Tensor>) = if nivel == 0 {
                    // prefixo entra junto na 1ª avaliação
                    let prefixo_ids = Tensor::zeros((b, N_PREFIXO), candle_core::DType::U32, dev)?;
                    let pe = g.emb_niveis.tabela.fwd(&prefixo_ids)?;
                    let e = Tensor::cat(&[pe, entrada_cond], 1)?;
                    let h = self.forward_bloco(g, &e, &mut caches_cond, nivel,
                        film_cond.as_ref().map(|(r, n)| (r, *n)), &mut estados_b0_cond)?;
                    self.nfe_contador += 1;
                    let hc = h.narrow(1, N_PREFIXO, n_l)?;
                    let hu = if let Some(eu) = entrada_uncond {
                        let eu2 = Tensor::cat(&[{
                            let prefixo_ids2 = Tensor::zeros((b, N_PREFIXO), candle_core::DType::U32, dev)?;
                            g.emb_niveis.tabela.fwd(&prefixo_ids2)?
                        }, eu], 1)?;
                        let hu = self.forward_bloco(g, &eu2, &mut caches_uncond, nivel,
                            film_uncond.as_ref().map(|(r, n)| (r, *n)), &mut estados_b0_uncond)?;
                        self.nfe_contador += 1;
                        Some(hu.narrow(1, N_PREFIXO, n_l)?)
                    } else { None };
                    (hc, hu)
                } else {
                    let h = self.forward_bloco(g, &entrada_cond, &mut caches_cond, nivel,
                        film_cond.as_ref().map(|(r, n)| (r, *n)), &mut estados_b0_cond)?;
                    self.nfe_contador += 1;
                    let hu = if let Some(eu) = entrada_uncond {
                        let hu = self.forward_bloco(g, &eu, &mut caches_uncond, nivel,
                            film_uncond.as_ref().map(|(r, n)| (r, *n)), &mut estados_b0_uncond)?;
                        self.nfe_contador += 1;
                        Some(hu)
                    } else { None };
                    (h, hu)
                };
                let ids = self.amostrar_nivel(g, &h_c, h_u.as_ref(), nivel, rng, greedy)?;
                codigos.push(ids.clone());
                zhats = t.reconstruir_de_codigos(&codigos, nivel + 1, b, dev)?;
            }
            Ok((codigos, self.nfe_contador - nfe_inicio))
        }
    }

    unsafe fn forward_bloco(
        &mut self,
        g: &mut GeradorConceitos,
        entrada: &Tensor,
        caches: &mut Vec<CacheKv>,
        nivel: usize,
        film: Option<(&Tensor, usize)>,
        estados_b0: &mut Tensor,
    ) -> Result<Tensor> {
        let n_novo = entrada.dim(1)?;
        let n_cache = caches.first().and_then(|c| c.k.as_ref()).map(|k| k.dims()[2]).unwrap_or(0);
        let total_vis = n_cache + n_novo;
        // block-causal incremental: posições novas veem cache + próprio bloco
        let m = vec![0f32; n_novo * total_vis];
        let mascara = Tensor::from_vec(m, (n_novo, total_vis), entrada.device())?;
        let niveis_novos = vec![nivel + 1; n_novo];
        let mut niveis_totais = vec![0usize; N_PREFIXO.min(n_cache)];
        let restante = n_cache.saturating_sub(N_PREFIXO);
        let mut usados = 0usize;
        for (li, &nl) in g.niveis.iter().enumerate().take(nivel) {
            if usados >= restante { break; }
            let tomar = nl.min(restante - usados);
            niveis_totais.extend(std::iter::repeat(li + 1).take(tomar));
            usados += tomar;
        }
        niveis_totais.extend(niveis_novos.clone());
        // estados pós-bloco-0 das posições novas (PRÉ-FiLM — como forward_treino)
        let mut novos_b0: Tensor = Tensor::zeros((0, 0, 0), candle_core::DType::F32, entrada.device())?;
        let h = g.tronco.forward_nivel(entrada, caches, &mascara, &niveis_novos, &niveis_totais,
            film, Some(&mut novos_b0))?;
        // acumula para o âncora dos níveis seguintes
        *estados_b0 = if estados_b0.dims()[1] == 0 {
            novos_b0
        } else {
            Tensor::cat(&[estados_b0.clone(), novos_b0], 1)?
        };
        Ok(h)
    }

    unsafe fn amostrar_nivel(
        &mut self,
        g: &mut GeradorConceitos,
        h_cond: &Tensor,
        h_uncond: Option<&Tensor>,
        nivel: usize,
        rng: &mut StdRng,
        greedy: bool,
    ) -> Result<Tensor> {
        let logits_cond = g.cabeca.cabecas[nivel].fwd(h_cond)?;
        let plano = match h_uncond {
            Some(hu) => {
                let lu = g.cabeca.cabecas[nivel].fwd(hu)?;
                // interpolação de logits (D.2(iii)): l̃ = l_u + s·(l_c − l_u)
                let dif = (&logits_cond - &lu)?.affine(self.cfg_scale, 0.0)?;
                (&lu + &dif)?
            }
            None => logits_cond,
        };
        let (b, n_l, _) = (plano.dim(0)?, plano.dim(1)?, plano.dim(2)?);
        let plano2 = plano.reshape((b * n_l, g.k))?;
        // temperatura-auto por nível (D.3)
        let tau = self.calibrar_tau(&plano2, 8)?;
        self.temperaturas[nivel] = tau;
        let probs = candle_nn::ops::softmax_last_dim(&plano2.affine(1.0 / tau.max(1e-3), 0.0)?)?;
        let plana = probs.flatten_all()?.to_vec1::<f32>()?;
        let k = g.k;
        let mut ids = Vec::with_capacity(b * n_l);
        for i in 0..b * n_l {
            if greedy {
                let (mut melhor, mut mi) = (-1f32, 0usize);
                for (j, v) in plana[i * k..(i + 1) * k].iter().enumerate() {
                    if *v > melhor { melhor = *v; mi = j; }
                }
                ids.push(mi as u32);
            } else {
                let r: f32 = rng.gen();
                let mut acc = 0f32;
                let mut esc = k - 1;
                for (j, v) in plana[i * k..(i + 1) * k].iter().enumerate() {
                    acc += *v;
                    if r <= acc { esc = j; break; }
                }
                ids.push(esc as u32);
            }
        }
        Tensor::from_vec(ids, (b, n_l), plano.device())
    }
}

/// H(softmax(logits/τ))/log K ∈ [0,1] — monotônica crescente em τ.
fn entropia_normalizada_de_logits(logits: &Tensor, tau: f64) -> Result<f64> {
    let escalados = logits.affine(1.0 / tau.max(1e-3), 0.0)?;
    let probs = candle_nn::ops::softmax_last_dim(&escalados)?;
    let p = probs.clamp(1e-9f64, 1.0)?;
    let h = (&p * &p.log()?)?.affine(-1.0, 0.0)?.sum(D::Minus1)?.mean_all()?.to_scalar::<f32>()? as f64;
    let logk = (logits.dim(D::Minus1)? as f64).ln();
    Ok(h / logk.max(1e-9))
}

/// MBR (Teo C.4/T6): escolhe o candidato com maior utilidade média contra os
/// demais. Utilidade = concordância de códigos por nível (proxy em [0,1],
/// limitada — satisfaz a hipótese do C.4; a utilidade de ensemble completa do
/// projeto é ponto de extensão via `ensemble/votacao.py`).
pub fn mbr_escolher(candidatos: &[Vec<Tensor>]) -> usize {
    let n = candidatos.len();
    if n <= 1 { return 0; }
    let mut melhor = 0usize;
    let mut melhor_u = -1f64;
    for j in 0..n {
        let mut u = 0f64;
        let mut pares = 0usize;
        for k in 0..n {
            if k == j { continue; }
            for nivel in 0..candidatos[j].len() {
                let a = candidatos[j][nivel].flatten_all().unwrap().to_vec1::<u32>().unwrap();
                let b = candidatos[k][nivel].flatten_all().unwrap().to_vec1::<u32>().unwrap();
                let iguais = a.iter().zip(b.iter()).filter(|(x, y)| x == y).count();
                u += iguais as f64 / a.len().max(1) as f64;
                pares += 1;
            }
        }
        let u = u / pares.max(1) as f64;
        if u > melhor_u { melhor_u = u; melhor = j; }
    }
    melhor
}