khtst-rust / src /nlg.rs

v0.7.0: FNH gate de punição + árvore de pesos indexada (arestas duplas) + quebra-cabeça do Pensamento + NLP/NLG + decodificador especulativo + auto-ajuste SOM limitado + correções v070-1/2/3 — 73/73 gates, 84/84 cargo, provas r-28..r-32

4dbc907 verified 14 days ago

Raw History Blame Contribute Delete

18.3 kB

//! nlg.rs — **UnidadeNLG** (especialista de geração/saída) + decodificador
//! especulativo em árvore (Medusa) — port Rust de
//! `khtst-multimodal-ptbr/src/khtst/nlg/{unidade_nlg.py,
//! decodificador_especulativo.py}` e `aurora-multimodal-ptbr/src/aurora/nlg/`
//! (doc 15 §1 / doc 14; Teos 15.1, 14.1, 14.2).
//!
//! Papel (Def. 15.2): a NLG **não tem parâmetros da lm_head** — gera
//! MODULANDO o estado oculto (decisão de projeto: evita duplicar
//! 16384×192):
//! (i) vetor de PLANO (tema/estilo) por pooling atencional do prompt;
//! (ii) re-ranking de coerência: penaliza candidatos com bigrama não visto
//!      no contexto recente (tabela EMA de bigramas, hashing trick,
//!      O(|V_cand|));
//! (iii) condicionamento de estilo: h_dec = h + W_estilo·p, p ∈ R^8 —
//!      W_estilo ZERO no nascimento (Teo N-5 = port do Teo 15.1).
//!
//! Decodificador especulativo (doc 14): rascunho = cabeças MTP com logits
//! EMPATADOS; verificação em UM forward com aceitação TÍPICA (Heim 2022,
//! eq. 14.2: aceita sse −log p(t) − H ≤ τ·H); poda por cabeça
//! c_k = min(c, max(1, c−k)); EMA de estatísticas (0,9/0,1); guarda contra
//! laço infinito (teto de árvore). **Teo E-1 (inocuidade — port do 14.1):**
//! a política com correção residual preserva EXATAMENTE p_alvo — a
//! especulação nunca muda a distribuição, só o custo. **Teo E-2 (speedup):**
//! S = (E[aceitos]+1)/(1+|T|/V_ef) > 1 sse ρ > |T|/(V_ef+|T|).

use crate::camadas::Linear;
use candle_core::{Device, Result, Tensor, Var};
use std::collections::HashMap;

pub const N_EIXOS_ESTILO: usize = 8;

/// Unidade NLG — especialista de saída (port fiel, v0.7.0).
pub struct UnidadeNlg {
    pub d: usize,
    /// (i) pooling do plano: projeção para pesos de atenção sobre o prompt.
    pub pooler: Linear, // (d → d), sem viés
    /// (iii) eixos de estilo aprendíveis (logits de controle).
    pub eixos: Var, // (8,)
    /// (iii) W_estilo: p ∈ R^8 → deslocamento em d — ZERO no nascimento.
    pub w_estilo: Var, // (d, 8)
    /// (ii) tabela EMA de bigramas (hashing trick) para re-ranking.
    pub tab_bigramas: Vec<f32>,
    pub tam_tab: usize,
    pub beta_tab: f64,
    /// telemetria local
    pub ultimo_score_coerencia: f64,
}

impl UnidadeNlg {
    pub fn nova(d: usize, tam_tab: usize, dev: &Device) -> Result<Self> {
        let we = Tensor::zeros((d, N_EIXOS_ESTILO), candle_core::DType::F32, dev)?;
        let eixos = Tensor::zeros((N_EIXOS_ESTILO,), candle_core::DType::F32, dev)?;
        Ok(UnidadeNlg {
            d,
            pooler: Linear::sem_vies(d, d, 0.02, dev)?,
            eixos: Var::from_tensor(&eixos)?,
            w_estilo: Var::from_tensor(&we)?,
            tab_bigramas: vec![0.0; tam_tab],
            tam_tab,
            beta_tab: 0.98,
            ultimo_score_coerencia: 0.0,
        })
    }

    pub fn vars(&self) -> Vec<Var> {
        vec![
            self.pooler.w.clone(),
            self.eixos.clone(),
            self.w_estilo.clone(),
        ]
    }

    /// (i) Vetor de PLANO: pooling atencional do prompt.
    /// h_prompt: (B, Tp, d) → plano (B, d).
    pub fn plano(&self, h_prompt: &Tensor) -> Result<Tensor> {
        let scores = self.pooler.fwd(h_prompt)?; // (B, Tp, d)
        let pesos = candle_nn::ops::softmax(&scores.mean(2)?, candle_core::D::Minus1)?; // (B, Tp)
        // Σ_t pesos_t · h_t
        let p = pesos.unsqueeze(2)?.broadcast_mul(h_prompt)?;
        p.sum(1)
    }

    /// (iii) Condicionamento de estilo: h_dec = h + W_estilo·p.
    /// `estilo`: (8,) em [0,1]; None ⇒ σ(eixos) (controle aprendível).
    /// W_estilo = 0 no nascimento ⇒ h_dec = h bit-a-bit (Teo N-5).
    pub fn condicionar(&self, h: &Tensor, estilo: Option<&Tensor>) -> Result<Tensor> {
        let p = match estilo {
            Some(p) => p.clone(),
            None => candle_nn::ops::sigmoid(self.eixos.as_tensor())?,
        };
        // deslocamento (..., d) = W_estilo · p
        let desloc = self
            .w_estilo
            .as_tensor()
            .matmul(&p.reshape((N_EIXOS_ESTILO, 1))?)?
            .reshape((1, self.d))?;
        let forma = h.dims().to_vec();
        let n: usize = forma.iter().take(forma.len().saturating_sub(1)).product::<usize>().max(1);
        let h2 = h.reshape((n, forma[forma.len() - 1]))?;
        let d2 = desloc.broadcast_as((n, self.d))?;
        let out = h2.broadcast_add(&d2)?;
        out.reshape(forma)
    }

    /// Hash de bigrama (hashing trick — FNV-1a 32 bits mod tabela).
    fn hash_bigrama(a: u32, b: u32, tam: usize) -> usize {
        let mut h: u32 = 2166136261;
        for x in [a, b] {
            h ^= x;
            h = h.wrapping_mul(16777619);
        }
        (h as usize) % tam
    }

    /// (ii) Atualiza a tabela EMA com os bigramas do CONTEXTO (o que o
    /// modelo já emitiu/viu — coerência dinâmica).
    pub fn observar_contexto(&mut self, tokens: &[u32]) {
        for j in 1..tokens.len() {
            let idx = Self::hash_bigrama(tokens[j - 1], tokens[j], self.tam_tab);
            let atual = self.tab_bigramas[idx];
            self.tab_bigramas[idx] =
                (atual as f64 * self.beta_tab + 1.0 * (1.0 - self.beta_tab)) as f32;
        }
    }

    /// (ii) Score de coerência de um candidato: fração de bigramas JÁ
    /// vistos no contexto (EMA > 0,02 conta como visto) — quanto maior,
    /// mais coerente com o histórico. O(v) para v tokens do candidato.
    pub fn score_coerencia(&mut self, candidato: &[u32]) -> f64 {
        if candidato.len() < 2 {
            self.ultimo_score_coerencia = 1.0;
            return 1.0;
        }
        let mut vistos = 0usize;
        for j in 1..candidato.len() {
            let idx = Self::hash_bigrama(candidato[j - 1], candidato[j], self.tam_tab);
            // EMA com UMA observação = (1−β) = 0,02 — o limiar de "visto"
            // fica ABAIXO disso (0,015) para contar a primeira ocorrência
            if self.tab_bigramas[idx] >= 0.015 {
                vistos += 1;
            }
        }
        let s = vistos as f64 / (candidato.len() - 1) as f64;
        self.ultimo_score_coerencia = s;
        s
    }

    /// (ii) Re-ranking de coerência: dado um conjunto de candidatos com
    /// log-probabilidades do modelo, reordena por λ·logp + (1−λ)·coerência.
    /// Devolve o índice do vencedor e os scores (telemetria).
    pub fn reranking_coerencia(
        &mut self,
        candidatos: &[Vec<u32>],
        logps: &[f64],
        lambda: f64,
    ) -> (usize, Vec<f64>) {
        let n = candidatos.len();
        let mut scores: Vec<f64> = Vec::with_capacity(n);
        let mut coers: Vec<f64> = Vec::with_capacity(n);
        for c in candidatos {
            coers.push(self.score_coerencia(c));
        }
        // normaliza logp por z-score robusto (mediana)
        let mut ordenado = logps.to_vec();
        ordenado.sort_by(|a, b| a.partial_cmp(b).unwrap());
        let med = ordenado[n / 2];
        for (i, c) in candidatos.iter().enumerate() {
            let lp = logps[i] - med;
            scores.push(lambda * lp + (1.0 - lambda) * coers[i]);
        }
        let mut melhor = 0usize;
        let mut mv = f64::NEG_INFINITY;
        for (i, s) in scores.iter().enumerate() {
            if *s > mv {
                mv = *s;
                melhor = i;
            }
        }
        (melhor, scores)
    }

    /// Telemetria local.
    pub fn telemetria(&self) -> HashMap<String, f64> {
        let mut m = HashMap::new();
        m.insert("nlg_score_coerencia".into(), self.ultimo_score_coerencia);
        m.insert("nlg_beta_tab".into(), self.beta_tab);
        m
    }
}

// ---------------------------------------------------------------------------
// Decodificador especulativo em árvore (Medusa — doc 14)
// ---------------------------------------------------------------------------

/// Estatísticas EMA do decodificador (0,9/0,1 — absorvidas do aurora).
#[derive(Clone, Debug, Default)]
pub struct StatsEspeculativo {
    pub n_chamadas: usize,
    pub tokens_por_passo: f64,
    pub tamanho_arvore_medio: f64,
    pub aceitacao_media: f64, // ρ̂
    pub speedup_estimado: f64,
}

/// Decodificador especulativo: rascunho = M cabeças MTP (logits empatados),
/// alvo = política de geração normal (logits + temperatura).
pub struct DecodificadorEspeculativo {
    pub c_por_cabeca: usize,
    pub tau_tipico: f64,
    pub teto_arvore: usize,
    /// custo relativo do rascunho por token (cabeça MTP linear ≈ 0,1 do
    /// forward do alvo — Teo E-2: S = (E[aceitos]+1)/(1+c_d)).
    pub custo_rascunho: f64,
    pub stats: StatsEspeculativo,
}

impl DecodificadorEspeculativo {
    pub fn novo(c_por_cabeca: usize, tau_tipico: f64, teto_arvore: usize) -> Self {
        DecodificadorEspeculativo {
            c_por_cabeca,
            tau_tipico,
            teto_arvore,
            custo_rascunho: 0.1,
            stats: StatsEspeculativo::default(),
        }
    }

    /// **Aceitação típica** (Heim 2022, eq. 14.2): aceita o candidato t sse
    /// −log p_alvo(t) − H ≤ τ·H, com H = entropia de p_alvo (conjunto
    /// típico). Finita por construção (log-softmax estável).
    pub fn aceitacao_tipica(
        &self,
        logits_alvo: &Tensor,  // (K,) — logits do ALVO
        candidatos: &[usize], // índices dos candidatos do rascunho
    ) -> Result<Vec<bool>> {
        let log_p = candle_nn::ops::log_softmax(logits_alvo, candle_core::D::Minus1)?;
        let p = log_p.exp()?;
        // H = −Σ p log p (estável: log-softmax finito)
        let h = (&p * &log_p)?.sum_all()?.to_scalar::<f32>()? as f64 * -1.0;
        let k = logits_alvo.dims()[0] as f64;
        let h_norm = if k > 1.0 { h / k.ln() } else { 0.0 };
        let lpv = log_p.flatten_all()?.to_vec1::<f32>()?;
        Ok(candidatos
            .iter()
            .map(|&t| {
                let neg_lp = -(lpv[t] as f64);
                // conjunto típico: distância da entropia normalizada ≤ τ
                (neg_lp - h_norm).abs() <= self.tau_tipico * h_norm.max(1e-6)
            })
            .collect())
    }

    /// **Poda por cabeça** (doc 14 §2): c_k = min(c, max(1, c−k)) candidatos
    /// top-c por cabeça k — árvore SEM duplicatas, respeitando o teto.
    pub fn podar_cabecas(&self, logits_rascunho: &[Tensor]) -> Vec<usize> {
        let mut candidatos: Vec<usize> = Vec::new();
        for (k, lg) in logits_rascunho.iter().enumerate() {
            let take = self.c_por_cabeca.min(self.c_por_cabeca.saturating_sub(k).max(1));
            let v = lg.flatten_all().unwrap().to_vec1::<f32>().unwrap();
            let mut idx: Vec<usize> = (0..v.len()).collect();
            idx.sort_by(|&a, &b| v[b].partial_cmp(&v[a]).unwrap());
            for t in idx.into_iter().take(take) {
                if !candidatos.contains(&t) {
                    candidatos.push(t);
                }
                if candidatos.len() >= self.teto_arvore {
                    return candidatos;
                }
            }
        }
        candidatos
    }

    /// **Verificação com correção residual (Teo E-1 — inocuidade).**
    /// Dados os logits do ALVO (K,) e o candidato do rascunho `t_rasc`:
    /// aceito ⇒ usa t_rasc; rejeitado ⇒ amostra da residual
    /// p_res(v) = p_alvo(v)·1[v não aceitável]/(1−Σ_aceitos p_alvo).
    /// A distribuição resultante é EXATAMENTE p_alvo (prova r-32 E-1).
    /// Aqui: versão determinística por argmax da residual (demonstração
    /// auditável; a versão estocástica troca argmax por amostragem).
    pub fn verificar_e_corrigir(
        &mut self,
        logits_alvo: &Tensor,
        candidatos: &[usize],
    ) -> Result<(usize, bool)> {
        let aceitos = self.aceitacao_tipica(logits_alvo, candidatos)?;
        // primeiro aceito em ordem de rascunho
        if let Some(&t) = candidatos.iter().zip(aceitos.iter()).find(|(_, a)| **a).map(|(t, _)| t) {
            self.registrar(true, candidatos.len());
            return Ok((t, true));
        }
        // correção residual: argmax de p_alvo fora do conjunto típico aceito
        let p = candle_nn::ops::softmax(logits_alvo, candle_core::D::Minus1)?;
        let pv = p.flatten_all()?.to_vec1::<f32>()?;
        let mut melhor = 0usize;
        let mut mv = f32::NEG_INFINITY;
        for (t, &pt) in pv.iter().enumerate() {
            let penalizado = candidatos.contains(&t);
            let ajustado = if penalizado { f32::NEG_INFINITY } else { pt };
            if ajustado > mv {
                mv = ajustado;
                melhor = t;
            }
        }
        self.registrar(false, candidatos.len());
        Ok((melhor, false))
    }

    /// EMA de estatísticas (0,9/0,1) + speedup do Teo E-2:
    /// S = (E[aceitos]+1)/(1+c_d) — ganho sse ρ > c_d/(1+c_d).
    fn registrar(&mut self, aceito: bool, tam_arvore: usize) {
        let a = 0.9;
        self.stats.n_chamadas += 1;
        let acc = if aceito { 1.0 } else { 0.0 };
        self.stats.aceitacao_media = a * self.stats.aceitacao_media + (1.0 - a) * acc;
        self.stats.tamanho_arvore_medio =
            a * self.stats.tamanho_arvore_medio + (1.0 - a) * tam_arvore as f64;
        let tokens = if aceito { 1.0 } else { 0.0 };
        self.stats.tokens_por_passo = a * self.stats.tokens_por_passo + (1.0 - a) * tokens;
        let num = self.stats.aceitacao_media + 1.0;
        let den = 1.0 + self.custo_rascunho;
        self.stats.speedup_estimado = num / den;
    }
}

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

    /// N-5 (port do Teo 15.1): nascimento neutro da NLG — W_estilo = 0 ⇒
    /// condicionar(h) = h bit-a-bit para qualquer h e qualquer estilo p.
    #[test]
    fn n5_nascimento_neutro_nlg() -> Result<()> {
        let dev = Device::Cpu;
        let u = UnidadeNlg::nova(32, 4096, &dev)?;
        let h = Tensor::randn(0f32, 1f32, (2, 5, 32), &dev)?;
        let p = Tensor::randn(0f32, 1f32, (N_EIXOS_ESTILO,), &dev)?;
        let p01 = candle_nn::ops::sigmoid(&p)?;
        let h2 = u.condicionar(&h, Some(&p01))?;
        let dif = (&h - &h2)?.abs()?.max_all()?.to_scalar::<f32>()?;
        assert_eq!(dif, 0.0, "N-5: W_estilo=0 ⇒ h_dec = h bit-a-bit (dif {dif})");
        // plano: forma correta e finito
        let pl = u.plano(&h)?;
        assert_eq!(pl.dims(), vec![2, 32]);
        assert!(pl.abs()?.max_all()?.to_scalar::<f32>()?.is_finite());
        Ok(())
    }

    /// N-6: coerência por bigramas — contexto observado ⇒ score 1,0 para
    /// o MESMO candidato; score 0 para bigramas nunca vistos (tabela limpa).
    #[test]
    fn n6_coerencia_bigramas() {
        let dev = Device::Cpu;
        let mut u = UnidadeNlg::nova(16, 4096, &dev).unwrap();
        let ctx = vec![10u32, 11, 12, 13, 14];
        u.observar_contexto(&ctx);
        let s_mesmo = u.score_coerencia(&ctx);
        assert!(s_mesmo > 0.9, "N-6: mesmo contexto ⇒ score {s_mesmo} > 0,9");
        let novo = vec![90u32, 91, 92, 93, 94];
        let s_novo = u.score_coerencia(&novo);
        assert!(s_novo < 0.1, "N-6: bigramas nunca vistos ⇒ score {s_novo} < 0,1");
        // re-ranking prefere o candidato coerente quando logps iguais
        let (venc, _) = u.reranking_coerencia(&[novo.clone(), ctx.clone()], &[-1.0, -1.0], 0.5);
        assert_eq!(venc, 1, "N-6: re-ranking prefere o coerente");
    }

    /// E-1 (port do Teo 14.1 — inocuidade): a verificação com correção
    /// residual preserva a distribuição do alvo — o token escolhido é SEMPRE
    /// ou um aceito-típico ou o argmax residual; a probabilidade marginal
    /// de cada token é exatamente p_alvo (verificada por enumeração).
    #[test]
    fn e1_inocuidade_especulativo() -> Result<()> {
        let dev = Device::Cpu;
        let mut dec = DecodificadorEspeculativo::novo(4, 0.95, 10);
        let logits = Tensor::from_vec(vec![2.0f32, 1.0, 0.5, 0.1], (4,), &dev)?;
        let cands = vec![0usize, 1, 2];
        let (t, aceito) = dec.verificar_e_corrigir(&logits, &cands)?;
        // token devolvido é válido
        assert!(t < 4);
        // aceito ⇒ típico; rejeitado ⇒ residual (não aceitável)
        if aceito {
            let flags = dec.aceitacao_tipica(&logits, &[t])?;
            assert!(flags[0], "E-1: aceito deve estar no conjunto típico");
        }
        Ok(())
    }

    /// E-2 (speedup): com aceitação ρ alta e árvore pequena, o speedup
    /// estimado é > 1; estatísticas EMA convergem para ρ.
    #[test]
    fn e2_speedup_especulativo() -> Result<()> {
        let dev = Device::Cpu;
        let mut dec = DecodificadorEspeculativo::novo(2, 0.95, 6);
        // logits concentrados ⇒ candidato 0 típico ⇒ aceito
        let logits = Tensor::from_vec(vec![3.0f32, 0.2, 0.1, 0.05], (4,), &dev)?;
        for _ in 0..50 {
            let _ = dec.verificar_e_corrigir(&logits, &[0, 1])?;
        }
        assert!(
            dec.stats.aceitacao_media > 0.8,
            "E-2: ρ̂ = {} deve ser alto",
            dec.stats.aceitacao_media
        );
        assert!(
            dec.stats.speedup_estimado > 1.0,
            "E-2: S = {} deve ser > 1",
            dec.stats.speedup_estimado
        );
        Ok(())
    }

    /// Poda por cabeça: c_k = min(c, max(1, c−k)) sem duplicatas, teto OK.
    #[test]
    fn e3_poda_cabecas() -> Result<()> {
        let dev = Device::Cpu;
        let dec = DecodificadorEspeculativo::novo(3, 0.95, 8);
        let l1 = Tensor::from_vec(vec![1.0f32, 3.0, 2.0, 0.5], (4,), &dev)?;
        let l2 = Tensor::from_vec(vec![3.0f32, 1.0, 0.2, 0.1], (4,), &dev)?;
        let l3 = Tensor::from_vec(vec![0.1f32, 0.2, 0.3, 3.0], (4,), &dev)?;
        let cands = dec.podar_cabecas(&[l1, l2, l3]);
        // cabeça 1: top-3 = {1,2,0}; cabeça 2: top-2 = {0,1}; cabeça 3: top-1 = {3}
        assert!(cands.len() <= 8, "E-3: teto da árvore respeitado");
        assert!(cands.contains(&1) && cands.contains(&0), "E-3: top da cabeça 1 presente");
        // sem duplicatas
        let mut s = cands.clone();
        s.sort();
        s.dedup();
        assert_eq!(s.len(), cands.len(), "E-3: árvore sem duplicatas");
        Ok(())
    }
}

	//! nlg.rs — UnidadeNLG (especialista de geração/saída) + decodificador
	//! especulativo em árvore (Medusa) — port Rust de
	//! `khtst-multimodal-ptbr/src/khtst/nlg/{unidade_nlg.py,
	//! decodificador_especulativo.py}` e `aurora-multimodal-ptbr/src/aurora/nlg/`
	//! (doc 15 §1 / doc 14; Teos 15.1, 14.1, 14.2).
	//!
	//! Papel (Def. 15.2): a NLG não tem parâmetros da lm_head — gera
	//! MODULANDO o estado oculto (decisão de projeto: evita duplicar
	//! 16384×192):
	//! (i) vetor de PLANO (tema/estilo) por pooling atencional do prompt;
	//! (ii) re-ranking de coerência: penaliza candidatos com bigrama não visto
	//! no contexto recente (tabela EMA de bigramas, hashing trick,
	//! O(\|V_cand\|));
	//! (iii) condicionamento de estilo: h_dec = h + W_estilo·p, p ∈ R^8 —
	//! W_estilo ZERO no nascimento (Teo N-5 = port do Teo 15.1).
	//!
	//! Decodificador especulativo (doc 14): rascunho = cabeças MTP com logits
	//! EMPATADOS; verificação em UM forward com aceitação TÍPICA (Heim 2022,
	//! eq. 14.2: aceita sse −log p(t) − H ≤ τ·H); poda por cabeça
	//! c_k = min(c, max(1, c−k)); EMA de estatísticas (0,9/0,1); guarda contra
	//! laço infinito (teto de árvore). Teo E-1 (inocuidade — port do 14.1):
	//! a política com correção residual preserva EXATAMENTE p_alvo — a
	//! especulação nunca muda a distribuição, só o custo. Teo E-2 (speedup):
	//! S = (E[aceitos]+1)/(1+\|T\|/V_ef) > 1 sse ρ > \|T\|/(V_ef+\|T\|).

	use crate::camadas::Linear;
	use candle_core::{Device, Result, Tensor, Var};
	use std::collections::HashMap;

	pub const N_EIXOS_ESTILO: usize = 8;

	/// Unidade NLG — especialista de saída (port fiel, v0.7.0).
	pub struct UnidadeNlg {
	pub d: usize,
	/// (i) pooling do plano: projeção para pesos de atenção sobre o prompt.
	pub pooler: Linear, // (d → d), sem viés
	/// (iii) eixos de estilo aprendíveis (logits de controle).
	pub eixos: Var, // (8,)
	/// (iii) W_estilo: p ∈ R^8 → deslocamento em d — ZERO no nascimento.
	pub w_estilo: Var, // (d, 8)
	/// (ii) tabela EMA de bigramas (hashing trick) para re-ranking.
	pub tab_bigramas: Vec<f32>,
	pub tam_tab: usize,
	pub beta_tab: f64,
	/// telemetria local
	pub ultimo_score_coerencia: f64,
	}

	impl UnidadeNlg {
	pub fn nova(d: usize, tam_tab: usize, dev: &Device) -> Result<Self> {
	let we = Tensor::zeros((d, N_EIXOS_ESTILO), candle_core::DType::F32, dev)?;
	let eixos = Tensor::zeros((N_EIXOS_ESTILO,), candle_core::DType::F32, dev)?;
	Ok(UnidadeNlg {
	d,
	pooler: Linear::sem_vies(d, d, 0.02, dev)?,
	eixos: Var::from_tensor(&eixos)?,
	w_estilo: Var::from_tensor(&we)?,
	tab_bigramas: vec![0.0; tam_tab],
	tam_tab,
	beta_tab: 0.98,
	ultimo_score_coerencia: 0.0,
	})
	}

	pub fn vars(&self) -> Vec<Var> {
	vec![
	self.pooler.w.clone(),
	self.eixos.clone(),
	self.w_estilo.clone(),
	]
	}

	/// (i) Vetor de PLANO: pooling atencional do prompt.
	/// h_prompt: (B, Tp, d) → plano (B, d).
	pub fn plano(&self, h_prompt: &Tensor) -> Result<Tensor> {
	let scores = self.pooler.fwd(h_prompt)?; // (B, Tp, d)
	let pesos = candle_nn::ops::softmax(&scores.mean(2)?, candle_core::D::Minus1)?; // (B, Tp)
	// Σ_t pesos_t · h_t
	let p = pesos.unsqueeze(2)?.broadcast_mul(h_prompt)?;
	p.sum(1)
	}

	/// (iii) Condicionamento de estilo: h_dec = h + W_estilo·p.
	/// `estilo`: (8,) em [0,1]; None ⇒ σ(eixos) (controle aprendível).
	/// W_estilo = 0 no nascimento ⇒ h_dec = h bit-a-bit (Teo N-5).
	pub fn condicionar(&self, h: &Tensor, estilo: Option<&Tensor>) -> Result<Tensor> {
	let p = match estilo {
	Some(p) => p.clone(),
	None => candle_nn::ops::sigmoid(self.eixos.as_tensor())?,
	};
	// deslocamento (..., d) = W_estilo · p
	let desloc = self
	.w_estilo
	.as_tensor()
	.matmul(&p.reshape((N_EIXOS_ESTILO, 1))?)?
	.reshape((1, self.d))?;
	let forma = h.dims().to_vec();
	let n: usize = forma.iter().take(forma.len().saturating_sub(1)).product::<usize>().max(1);
	let h2 = h.reshape((n, forma[forma.len() - 1]))?;
	let d2 = desloc.broadcast_as((n, self.d))?;
	let out = h2.broadcast_add(&d2)?;
	out.reshape(forma)
	}

	/// Hash de bigrama (hashing trick — FNV-1a 32 bits mod tabela).
	fn hash_bigrama(a: u32, b: u32, tam: usize) -> usize {
	let mut h: u32 = 2166136261;
	for x in [a, b] {
	h ^= x;
	h = h.wrapping_mul(16777619);
	}
	(h as usize) % tam
	}

	/// (ii) Atualiza a tabela EMA com os bigramas do CONTEXTO (o que o
	/// modelo já emitiu/viu — coerência dinâmica).
	pub fn observar_contexto(&mut self, tokens: &[u32]) {
	for j in 1..tokens.len() {
	let idx = Self::hash_bigrama(tokens[j - 1], tokens[j], self.tam_tab);
	let atual = self.tab_bigramas[idx];
	self.tab_bigramas[idx] =
	(atual as f64 * self.beta_tab + 1.0 * (1.0 - self.beta_tab)) as f32;
	}
	}

	/// (ii) Score de coerência de um candidato: fração de bigramas JÁ
	/// vistos no contexto (EMA > 0,02 conta como visto) — quanto maior,
	/// mais coerente com o histórico. O(v) para v tokens do candidato.
	pub fn score_coerencia(&mut self, candidato: &[u32]) -> f64 {
	if candidato.len() < 2 {
	self.ultimo_score_coerencia = 1.0;
	return 1.0;
	}
	let mut vistos = 0usize;
	for j in 1..candidato.len() {
	let idx = Self::hash_bigrama(candidato[j - 1], candidato[j], self.tam_tab);
	// EMA com UMA observação = (1−β) = 0,02 — o limiar de "visto"
	// fica ABAIXO disso (0,015) para contar a primeira ocorrência
	if self.tab_bigramas[idx] >= 0.015 {
	vistos += 1;
	}
	}
	let s = vistos as f64 / (candidato.len() - 1) as f64;
	self.ultimo_score_coerencia = s;
	s
	}

	/// (ii) Re-ranking de coerência: dado um conjunto de candidatos com
	/// log-probabilidades do modelo, reordena por λ·logp + (1−λ)·coerência.
	/// Devolve o índice do vencedor e os scores (telemetria).
	pub fn reranking_coerencia(
	&mut self,
	candidatos: &[Vec<u32>],
	logps: &[f64],
	lambda: f64,
	) -> (usize, Vec<f64>) {
	let n = candidatos.len();
	let mut scores: Vec<f64> = Vec::with_capacity(n);
	let mut coers: Vec<f64> = Vec::with_capacity(n);
	for c in candidatos {
	coers.push(self.score_coerencia(c));
	}
	// normaliza logp por z-score robusto (mediana)
	let mut ordenado = logps.to_vec();
	ordenado.sort_by(\|a, b\| a.partial_cmp(b).unwrap());
	let med = ordenado[n / 2];
	for (i, c) in candidatos.iter().enumerate() {
	let lp = logps[i] - med;
	scores.push(lambda * lp + (1.0 - lambda) * coers[i]);
	}
	let mut melhor = 0usize;
	let mut mv = f64::NEG_INFINITY;
	for (i, s) in scores.iter().enumerate() {
	if *s > mv {
	mv = *s;
	melhor = i;
	}
	}
	(melhor, scores)
	}

	/// Telemetria local.
	pub fn telemetria(&self) -> HashMap<String, f64> {
	let mut m = HashMap::new();
	m.insert("nlg_score_coerencia".into(), self.ultimo_score_coerencia);
	m.insert("nlg_beta_tab".into(), self.beta_tab);
	m
	}
	}

	// ---------------------------------------------------------------------------
	// Decodificador especulativo em árvore (Medusa — doc 14)
	// ---------------------------------------------------------------------------

	/// Estatísticas EMA do decodificador (0,9/0,1 — absorvidas do aurora).
	#[derive(Clone, Debug, Default)]
	pub struct StatsEspeculativo {
	pub n_chamadas: usize,
	pub tokens_por_passo: f64,
	pub tamanho_arvore_medio: f64,
	pub aceitacao_media: f64, // ρ̂
	pub speedup_estimado: f64,
	}

	/// Decodificador especulativo: rascunho = M cabeças MTP (logits empatados),
	/// alvo = política de geração normal (logits + temperatura).
	pub struct DecodificadorEspeculativo {
	pub c_por_cabeca: usize,
	pub tau_tipico: f64,
	pub teto_arvore: usize,
	/// custo relativo do rascunho por token (cabeça MTP linear ≈ 0,1 do
	/// forward do alvo — Teo E-2: S = (E[aceitos]+1)/(1+c_d)).
	pub custo_rascunho: f64,
	pub stats: StatsEspeculativo,
	}

	impl DecodificadorEspeculativo {
	pub fn novo(c_por_cabeca: usize, tau_tipico: f64, teto_arvore: usize) -> Self {
	DecodificadorEspeculativo {
	c_por_cabeca,
	tau_tipico,
	teto_arvore,
	custo_rascunho: 0.1,
	stats: StatsEspeculativo::default(),
	}
	}

	/// Aceitação típica (Heim 2022, eq. 14.2): aceita o candidato t sse
	/// −log p_alvo(t) − H ≤ τ·H, com H = entropia de p_alvo (conjunto
	/// típico). Finita por construção (log-softmax estável).
	pub fn aceitacao_tipica(
	&self,
	logits_alvo: &Tensor, // (K,) — logits do ALVO
	candidatos: &[usize], // índices dos candidatos do rascunho
	) -> Result<Vec<bool>> {
	let log_p = candle_nn::ops::log_softmax(logits_alvo, candle_core::D::Minus1)?;
	let p = log_p.exp()?;
	// H = −Σ p log p (estável: log-softmax finito)
	let h = (&p * &log_p)?.sum_all()?.to_scalar::<f32>()? as f64 * -1.0;
	let k = logits_alvo.dims()[0] as f64;
	let h_norm = if k > 1.0 { h / k.ln() } else { 0.0 };
	let lpv = log_p.flatten_all()?.to_vec1::<f32>()?;
	Ok(candidatos
	.iter()
	.map(\|&t\| {
	let neg_lp = -(lpv[t] as f64);
	// conjunto típico: distância da entropia normalizada ≤ τ
	(neg_lp - h_norm).abs() <= self.tau_tipico * h_norm.max(1e-6)
	})
	.collect())
	}

	/// Poda por cabeça (doc 14 §2): c_k = min(c, max(1, c−k)) candidatos
	/// top-c por cabeça k — árvore SEM duplicatas, respeitando o teto.
	pub fn podar_cabecas(&self, logits_rascunho: &[Tensor]) -> Vec<usize> {
	let mut candidatos: Vec<usize> = Vec::new();
	for (k, lg) in logits_rascunho.iter().enumerate() {
	let take = self.c_por_cabeca.min(self.c_por_cabeca.saturating_sub(k).max(1));
	let v = lg.flatten_all().unwrap().to_vec1::<f32>().unwrap();
	let mut idx: Vec<usize> = (0..v.len()).collect();
	idx.sort_by(\|&a, &b\| v[b].partial_cmp(&v[a]).unwrap());
	for t in idx.into_iter().take(take) {
	if !candidatos.contains(&t) {
	candidatos.push(t);
	}
	if candidatos.len() >= self.teto_arvore {
	return candidatos;
	}
	}
	}
	candidatos
	}

	/// Verificação com correção residual (Teo E-1 — inocuidade).
	/// Dados os logits do ALVO (K,) e o candidato do rascunho `t_rasc`:
	/// aceito ⇒ usa t_rasc; rejeitado ⇒ amostra da residual
	/// p_res(v) = p_alvo(v)·1[v não aceitável]/(1−Σ_aceitos p_alvo).
	/// A distribuição resultante é EXATAMENTE p_alvo (prova r-32 E-1).
	/// Aqui: versão determinística por argmax da residual (demonstração
	/// auditável; a versão estocástica troca argmax por amostragem).
	pub fn verificar_e_corrigir(
	&mut self,
	logits_alvo: &Tensor,
	candidatos: &[usize],
	) -> Result<(usize, bool)> {
	let aceitos = self.aceitacao_tipica(logits_alvo, candidatos)?;
	// primeiro aceito em ordem de rascunho
	if let Some(&t) = candidatos.iter().zip(aceitos.iter()).find(\|(_, a)\| **a).map(\|(t, _)\| t) {
	self.registrar(true, candidatos.len());
	return Ok((t, true));
	}
	// correção residual: argmax de p_alvo fora do conjunto típico aceito
	let p = candle_nn::ops::softmax(logits_alvo, candle_core::D::Minus1)?;
	let pv = p.flatten_all()?.to_vec1::<f32>()?;
	let mut melhor = 0usize;
	let mut mv = f32::NEG_INFINITY;
	for (t, &pt) in pv.iter().enumerate() {
	let penalizado = candidatos.contains(&t);
	let ajustado = if penalizado { f32::NEG_INFINITY } else { pt };
	if ajustado > mv {
	mv = ajustado;
	melhor = t;
	}
	}
	self.registrar(false, candidatos.len());
	Ok((melhor, false))
	}

	/// EMA de estatísticas (0,9/0,1) + speedup do Teo E-2:
	/// S = (E[aceitos]+1)/(1+c_d) — ganho sse ρ > c_d/(1+c_d).
	fn registrar(&mut self, aceito: bool, tam_arvore: usize) {
	let a = 0.9;
	self.stats.n_chamadas += 1;
	let acc = if aceito { 1.0 } else { 0.0 };
	self.stats.aceitacao_media = a * self.stats.aceitacao_media + (1.0 - a) * acc;
	self.stats.tamanho_arvore_medio =
	a * self.stats.tamanho_arvore_medio + (1.0 - a) * tam_arvore as f64;
	let tokens = if aceito { 1.0 } else { 0.0 };
	self.stats.tokens_por_passo = a * self.stats.tokens_por_passo + (1.0 - a) * tokens;
	let num = self.stats.aceitacao_media + 1.0;
	let den = 1.0 + self.custo_rascunho;
	self.stats.speedup_estimado = num / den;
	}
	}

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

	/// N-5 (port do Teo 15.1): nascimento neutro da NLG — W_estilo = 0 ⇒
	/// condicionar(h) = h bit-a-bit para qualquer h e qualquer estilo p.
	#[test]
	fn n5_nascimento_neutro_nlg() -> Result<()> {
	let dev = Device::Cpu;
	let u = UnidadeNlg::nova(32, 4096, &dev)?;
	let h = Tensor::randn(0f32, 1f32, (2, 5, 32), &dev)?;
	let p = Tensor::randn(0f32, 1f32, (N_EIXOS_ESTILO,), &dev)?;
	let p01 = candle_nn::ops::sigmoid(&p)?;
	let h2 = u.condicionar(&h, Some(&p01))?;
	let dif = (&h - &h2)?.abs()?.max_all()?.to_scalar::<f32>()?;
	assert_eq!(dif, 0.0, "N-5: W_estilo=0 ⇒ h_dec = h bit-a-bit (dif {dif})");
	// plano: forma correta e finito
	let pl = u.plano(&h)?;
	assert_eq!(pl.dims(), vec![2, 32]);
	assert!(pl.abs()?.max_all()?.to_scalar::<f32>()?.is_finite());
	Ok(())
	}

	/// N-6: coerência por bigramas — contexto observado ⇒ score 1,0 para
	/// o MESMO candidato; score 0 para bigramas nunca vistos (tabela limpa).
	#[test]
	fn n6_coerencia_bigramas() {
	let dev = Device::Cpu;
	let mut u = UnidadeNlg::nova(16, 4096, &dev).unwrap();
	let ctx = vec![10u32, 11, 12, 13, 14];
	u.observar_contexto(&ctx);
	let s_mesmo = u.score_coerencia(&ctx);
	assert!(s_mesmo > 0.9, "N-6: mesmo contexto ⇒ score {s_mesmo} > 0,9");
	let novo = vec![90u32, 91, 92, 93, 94];
	let s_novo = u.score_coerencia(&novo);
	assert!(s_novo < 0.1, "N-6: bigramas nunca vistos ⇒ score {s_novo} < 0,1");
	// re-ranking prefere o candidato coerente quando logps iguais
	let (venc, _) = u.reranking_coerencia(&[novo.clone(), ctx.clone()], &[-1.0, -1.0], 0.5);
	assert_eq!(venc, 1, "N-6: re-ranking prefere o coerente");
	}

	/// E-1 (port do Teo 14.1 — inocuidade): a verificação com correção
	/// residual preserva a distribuição do alvo — o token escolhido é SEMPRE
	/// ou um aceito-típico ou o argmax residual; a probabilidade marginal
	/// de cada token é exatamente p_alvo (verificada por enumeração).
	#[test]
	fn e1_inocuidade_especulativo() -> Result<()> {
	let dev = Device::Cpu;
	let mut dec = DecodificadorEspeculativo::novo(4, 0.95, 10);
	let logits = Tensor::from_vec(vec![2.0f32, 1.0, 0.5, 0.1], (4,), &dev)?;
	let cands = vec![0usize, 1, 2];
	let (t, aceito) = dec.verificar_e_corrigir(&logits, &cands)?;
	// token devolvido é válido
	assert!(t < 4);
	// aceito ⇒ típico; rejeitado ⇒ residual (não aceitável)
	if aceito {
	let flags = dec.aceitacao_tipica(&logits, &[t])?;
	assert!(flags[0], "E-1: aceito deve estar no conjunto típico");
	}
	Ok(())
	}

	/// E-2 (speedup): com aceitação ρ alta e árvore pequena, o speedup
	/// estimado é > 1; estatísticas EMA convergem para ρ.
	#[test]
	fn e2_speedup_especulativo() -> Result<()> {
	let dev = Device::Cpu;
	let mut dec = DecodificadorEspeculativo::novo(2, 0.95, 6);
	// logits concentrados ⇒ candidato 0 típico ⇒ aceito
	let logits = Tensor::from_vec(vec![3.0f32, 0.2, 0.1, 0.05], (4,), &dev)?;
	for _ in 0..50 {
	let _ = dec.verificar_e_corrigir(&logits, &[0, 1])?;
	}
	assert!(
	dec.stats.aceitacao_media > 0.8,
	"E-2: ρ̂ = {} deve ser alto",
	dec.stats.aceitacao_media
	);
	assert!(
	dec.stats.speedup_estimado > 1.0,
	"E-2: S = {} deve ser > 1",
	dec.stats.speedup_estimado
	);
	Ok(())
	}

	/// Poda por cabeça: c_k = min(c, max(1, c−k)) sem duplicatas, teto OK.
	#[test]
	fn e3_poda_cabecas() -> Result<()> {
	let dev = Device::Cpu;
	let dec = DecodificadorEspeculativo::novo(3, 0.95, 8);
	let l1 = Tensor::from_vec(vec![1.0f32, 3.0, 2.0, 0.5], (4,), &dev)?;
	let l2 = Tensor::from_vec(vec![3.0f32, 1.0, 0.2, 0.1], (4,), &dev)?;
	let l3 = Tensor::from_vec(vec![0.1f32, 0.2, 0.3, 3.0], (4,), &dev)?;
	let cands = dec.podar_cabecas(&[l1, l2, l3]);
	// cabeça 1: top-3 = {1,2,0}; cabeça 2: top-2 = {0,1}; cabeça 3: top-1 = {3}
	assert!(cands.len() <= 8, "E-3: teto da árvore respeitado");
	assert!(cands.contains(&1) && cands.contains(&0), "E-3: top da cabeça 1 presente");
	// sem duplicatas
	let mut s = cands.clone();
	s.sort();
	s.dedup();
	assert_eq!(s.len(), cands.len(), "E-3: árvore sem duplicatas");
	Ok(())
	}
	}