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 | //! 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 | |
| } | |