v0.9.0: análise/prova/teste das críticas Qwen3.8 lote 2 (r-34) — SIMD com despacho em runtime bit-a-bit (X-1), oráculo candle_nn::optim::AdamW (nosso 1,53× mais rápido), Thompson Sampling (regret 0,0019 vs UCB1 0,0156), GNG por entropia com piso (T-G1'), JEPA de fração adaptativa (bit-a-bit com 1 braço), configs tipadas, entrada numpy (2008× na fronteira), avaliação com dados reais — 104/104 gates, 111/111 cargo
d444358 verified | //! Tokenizador de conceitos NCM (Estágio 1 / fase A) — port de | |
| //! `conceitos/tokenizador_conceitos.py` (n-04 §2.3) com a síntese GSQ+RCO | |
| //! (Teo D.4): codebooks Stiefel por nível, seleção GSQ (gradiente exato), | |
| //! nested dropout (Teo A.2), commitment (Teo A.0) e reconstrução por CE | |
| //! token-level no vocabulário. | |
| //! | |
| //! Recursão residual (n-02 §0): | |
| //! z_1 = Interp_{n_1}(E(x)); r_ℓ = Quant_{C_ℓ}(z_ℓ); | |
| //! z_{ℓ+1} = Interp_{n_{ℓ+1}}(z_ℓ) − Interp_{n_{ℓ+1}}(r_ℓ); | |
| //! ẑ_ℓ = Σ_{j≤ℓ} Interp_{n_ℓ}(r_j) (reconstrução cumulativa) | |
| use candle_core::{Device, Result, Tensor, Var, D}; | |
| use crate::atencao::BlocoAdba; | |
| use crate::camadas::{Embedding, Linear, RmsNorm}; | |
| use crate::config::ConfigKhtst; | |
| use crate::mascaras::{interp_para, matriz_interp}; | |
| use crate::quantizador::{ModoQuant, QuantizadorNivel}; | |
| use rand::prelude::*; | |
| pub struct TokenizadorConceitos { | |
| // Encoder (treino apenas; atenção plena) | |
| pub enc_emb: Embedding, | |
| pub enc_blocos: Vec<BlocoAdba>, | |
| pub enc_norm: RmsNorm, | |
| // Projeções por nível d_modelo → d_lat + quantizadores Stiefel/GSQ | |
| pub proj_niveis: Vec<Linear>, | |
| pub quants: Vec<QuantizadorNivel>, | |
| // Decoder AR condicionado em ẑ (B.1: p_ψ(x|c) = Π_i p(x_i|x_{<i}, ẑ)) | |
| pub dec_proj: Linear, | |
| pub dec_emb: Embedding, | |
| pub dec_blocos: Vec<BlocoAdba>, | |
| pub dec_norm: RmsNorm, | |
| pub dec_head: Linear, | |
| // máscaras pré-computadas | |
| pub mascara_causal_dec: Tensor, | |
| pub niveis: Vec<usize>, | |
| pub d_lat: usize, | |
| pub vocab: usize, | |
| pub cfg_ncm_beta_commit: f64, | |
| pub modo: ModoQuant, | |
| } | |
| /// Saída da codificação: códigos por nível + reconstruções cumulativas. | |
| pub struct Codigos { | |
| pub ids: Vec<Tensor>, // u32 (b, n_ℓ) | |
| pub zhats: Vec<Tensor>, // (b, n_ℓ, d_lat) — ẑ_ℓ na granularidade do nível | |
| } | |
| impl TokenizadorConceitos { | |
| pub fn novo(cfg: &ConfigKhtst, dev: &Device) -> Result<Self> { | |
| let m = &cfg.modelo; | |
| let n = &cfg.ncm; | |
| let l = n.niveis.len(); | |
| let n_total_dec = m.comprimento_ctx; | |
| let mascara_causal_dec = causal_aditiva(n_total_dec, dev)?; | |
| let mut enc_blocos = Vec::new(); | |
| let mut dec_blocos = Vec::new(); | |
| for _ in 0..2 { | |
| enc_blocos.push(BlocoAdba::novo(m.d_modelo, m.n_cabecas, m.d_ff, l, false, false, false, dev)?); | |
| dec_blocos.push(BlocoAdba::novo(m.d_modelo, m.n_cabecas, m.d_ff, l, false, false, false, dev)?); | |
| } | |
| let modo = ModoQuant::de_str(&n.modo_quantizacao); | |
| let mut quants = Vec::with_capacity(l); | |
| for _ in 0..l { | |
| quants.push(QuantizadorNivel::novo(n.k_por_nivel, n.d_lat, modo, n.ema_codebook, &n.escala_r, dev)?); | |
| } | |
| let mut proj = Vec::with_capacity(l); | |
| for _ in 0..l { | |
| proj.push(Linear::sem_vies(m.d_modelo, n.d_lat, 0.02, dev)?); | |
| } | |
| Ok(TokenizadorConceitos { | |
| enc_emb: Embedding::novo(m.vocab, m.d_modelo, dev)?, | |
| enc_blocos, | |
| enc_norm: RmsNorm::novo(m.d_modelo, 1e-5, dev)?, | |
| proj_niveis: proj, | |
| quants, | |
| dec_proj: Linear::sem_vies(n.d_lat, m.d_modelo, 0.02, dev)?, | |
| dec_emb: Embedding::novo(m.vocab, m.d_modelo, dev)?, | |
| dec_blocos, | |
| dec_norm: RmsNorm::novo(m.d_modelo, 1e-5, dev)?, | |
| dec_head: Linear::sem_vies(m.d_modelo, m.vocab, 0.02, dev)?, | |
| mascara_causal_dec, | |
| niveis: n.niveis.clone(), | |
| d_lat: n.d_lat, | |
| vocab: m.vocab, | |
| cfg_ncm_beta_commit: n.beta_commit, | |
| modo, | |
| }) | |
| } | |
| /// Encoder: x (b, N) → z (b, N, d_modelo). Atenção plena (só treino). | |
| pub fn encoder(&self, ids: &Tensor) -> Result<Tensor> { | |
| let mut x = self.enc_emb.fwd(ids)?; | |
| for bloco in &self.enc_blocos { | |
| let mut cache = Default::default(); | |
| x = bloco.fwd(&x, None, &mut cache, None, None)?; | |
| } | |
| self.enc_norm.fwd(&x) | |
| } | |
| /// Codificação de TREINO (fase A): relaxamento GSQ suave diferenciável. | |
| /// Devolve (zhats_suaves, zs_por_nivel) para a perda e telemetria. | |
| fn codificar_treino(&mut self, ids: &Tensor) -> Result<(Vec<Tensor>, Vec<Tensor>)> { | |
| let dev = ids.device(); | |
| let z = self.encoder(ids)?; | |
| let mut zs = Vec::with_capacity(self.niveis.len()); | |
| let mut zhats = Vec::with_capacity(self.niveis.len()); | |
| // projeção ÚNICA d_modelo → d_lat na granularidade do nível 1 (n-02 §0); | |
| // a recursão residual vive INTEIRAMENTE em d_lat | |
| let n1 = self.niveis[0]; | |
| let mut z_corrente = self.proj_niveis[0].fwd(&interp_para(&z, n1, dev)?)?; | |
| let mut zhat_acum: Option<Tensor> = None; | |
| for (i, &n_l) in self.niveis.iter().enumerate() { | |
| let zp = z_corrente.clone(); | |
| let r_suave = self.quants[i].quantizar_treino(&zp)?; | |
| // telemetria + escalas + uso (sem gradiente) | |
| let (ids_duros, _valores) = self.quants[i].quantizar_duro(&zp)?; | |
| self.quants[i].registrar_uso(&ids_duros)?; | |
| self.quants[i].atualizar_escala(&zp)?; | |
| // reconstrução cumulativa na granularidade n_l | |
| let zhat_l = match &zhat_acum { | |
| Some(a) => (&interp_para(a, n_l, dev)? + &r_suave)?, | |
| None => r_suave.clone(), | |
| }; | |
| zhat_acum = Some(zhat_l.clone()); | |
| zhats.push(zhat_l); | |
| // resíduo na granularidade do próximo nível (Teo A.0/D.4) | |
| if i + 1 < self.niveis.len() { | |
| let n_prox = self.niveis[i + 1]; | |
| let res = (&zp - &r_suave)?; | |
| z_corrente = interp_para(&res, n_prox, dev)?; | |
| } | |
| zs.push(zp); | |
| } | |
| Ok((zhats, zs)) | |
| } | |
| /// Codificação CONGELADA (fase B / inferência): códigos DUROS + ẑ dos | |
| /// códigos (T determinístico — hipótese do Teo B.1). Recursão residual | |
| /// inteiramente em d_lat (projeção única — n-02 §0). | |
| pub fn codificar_congelado(&self, ids: &Tensor) -> Result<Codigos> { | |
| { | |
| let dev = ids.device(); | |
| let z = self.encoder(ids)?.detach(); | |
| let mut saida_ids = Vec::new(); | |
| let mut zhats = Vec::new(); | |
| let n1 = self.niveis[0]; | |
| let mut z_corrente = self.proj_niveis[0].fwd(&interp_para(&z, n1, dev)?)?; | |
| let mut zhat_acum: Option<Tensor> = None; | |
| for (i, &n_l) in self.niveis.iter().enumerate() { | |
| let zp = z_corrente.clone(); | |
| let (ids_l, valores_l) = self.quants[i].quantizar_duro(&zp)?; | |
| saida_ids.push(ids_l.reshape((zp.dims()[0], n_l))?); | |
| let zhat_l = match &zhat_acum { | |
| Some(a) => (&interp_para(a, n_l, dev)? + &valores_l)?, | |
| None => valores_l.clone(), | |
| }; | |
| zhat_acum = Some(zhat_l.clone()); | |
| zhats.push(zhat_l); | |
| if i + 1 < self.niveis.len() { | |
| let n_prox = self.niveis[i + 1]; | |
| let res = (&zp - &valores_l)?; | |
| z_corrente = interp_para(&res, n_prox, dev)?; | |
| } | |
| } | |
| Ok(Codigos { ids: saida_ids, zhats }) | |
| } | |
| } | |
| /// Reconstrói ẑ a partir de códigos (para corrupção κ e amostragem — | |
| /// o gerador nunca vê os vetores contínuos, apenas c (C.3)). | |
| pub fn reconstruir_de_codigos(&self, ids: &[Tensor], ate_nivel: usize, b: usize, dev: &Device) -> Result<Vec<Tensor>> { | |
| { | |
| let mut zhats = Vec::new(); | |
| let mut zhat_acum: Option<Tensor> = None; | |
| for (i, item) in self.niveis.iter().enumerate().take(ate_nivel.min(self.niveis.len())) { | |
| let plano = ids[i].reshape((b, self.niveis[i]))?; | |
| let valores = self.quants[i].materializar(&plano.flatten_all()?)?; | |
| let valores = valores.reshape((b, self.niveis[i], self.d_lat))?; | |
| let zhat_l = match &zhat_acum { | |
| Some(a) => (&interp_para(a, *item, dev)? + &valores)?, | |
| None => valores.clone(), | |
| }; | |
| zhat_acum = Some(zhat_l.clone()); | |
| zhats.push(zhat_l); | |
| } | |
| Ok(zhats) | |
| } | |
| } | |
| /// Decoder AR condicionado: ẑ (b, n_ℓ, d_lat) → logits (b, N, vocab). | |
| /// Entrada = proj(upsample_N(ẑ)) + emb(tokens deslocados). | |
| pub fn decoder(&self, zhat: &Tensor, alvo: Option<&Tensor>) -> Result<Tensor> { | |
| let dev = zhat.device(); | |
| let b = zhat.dim(0)?; | |
| let n_z = zhat.dim(1)?; | |
| let n_total = self.mascara_causal_dec.dim(0)?; | |
| let up = interp_para(zhat, n_total, dev)?; | |
| let cond = self.dec_proj.fwd(&up)?; | |
| let entrada_tokens: Tensor = match alvo { | |
| Some(t) => { | |
| // desloca à direita com BOS | |
| let bos = Tensor::full(crate::tokens::BOS as f32, (b, 1), dev)?.to_dtype(candle_core::DType::U32)?; | |
| let corpo = t.narrow(1, 0, t.dim(1)? - 1)?; | |
| Tensor::cat(&[&bos, &corpo], 1)? | |
| } | |
| None => Tensor::full(crate::tokens::BOS as f32, (b, n_total), dev)?.to_dtype(candle_core::DType::U32)?, | |
| }; | |
| let mut x = self.dec_emb.fwd(&entrada_tokens)?; | |
| x = (&x + &cond)?; | |
| for bloco in &self.dec_blocos { | |
| let mut cache = Default::default(); | |
| x = bloco.fwd(&x, Some(&self.mascara_causal_dec), &mut cache, None, None)?; | |
| } | |
| let x = self.dec_norm.fwd(&x)?; | |
| self.dec_head.fwd(&x) | |
| } | |
| /// Perda da fase A (n-04 §2.3): nested dropout + commitment. | |
| /// p(ℓ) uniforme (Teo A.2); RNG DEDICADO (pareamento — Teo 31.7). | |
| /// Devolve (tensor da perda total, métricas). | |
| pub fn perda_fase_a(&mut self, ids: &Tensor, rng: &mut StdRng) -> Result<(Tensor, PerdaFaseA)> { | |
| let (zhats, zs) = self.codificar_treino(ids)?; | |
| let l = zhats.len(); | |
| // sorteio do nível de parada (nested dropout) | |
| let nivel_parada = rng.gen_range(1..=l); | |
| let logits = self.decoder(&zhats[nivel_parada - 1], Some(ids))?; | |
| let (b, n_t, _) = (logits.dim(0)?, logits.dim(1)?, logits.dim(2)?); | |
| let plano = logits.reshape((b * n_t, self.vocab))?; | |
| let alvos = ids.reshape(b * n_t)?; | |
| let ce = candle_nn::loss::cross_entropy(&plano, &alvos)?; | |
| // commitment: β Σ_ℓ ‖sg[r_ℓ] − z_ℓ‖² (média por elemento) | |
| let mut commit_t: Option<Tensor> = None; | |
| let mut commit = 0f64; | |
| for (i, zp) in zs.iter().enumerate() { | |
| let (_ids, valores) = self.quants[i].quantizar_duro(zp)?; | |
| let dif = (&valores - zp)?; | |
| let termo = dif.sqr()?.mean_all()?; | |
| commit += termo.to_scalar::<f32>()? as f64; | |
| commit_t = Some(match commit_t { Some(s) => (&s + &termo)?, None => termo }); | |
| } | |
| let commit_t = commit_t.unwrap_or_else(|| ce.zeros_like().unwrap()); | |
| let total = (&ce + &commit_t.affine(self.cfg_ncm_beta_commit, 0.0)?)?; | |
| // precisão de reconstrução (telemetria) | |
| let preds = plano.argmax(D::Minus1)?; | |
| let acertos = preds.eq(&alvos)?.to_dtype(candle_core::DType::F32)?.sum_all()?.to_scalar::<f32>()? as f64; | |
| let acc = acertos / (b * n_t) as f64; | |
| // δ_ort por nível (V9) | |
| let deltas: Result<Vec<f64>> = self.quants.iter().map(|q| q.delta_ort()).collect(); | |
| let deltas = deltas?; | |
| for (i, d) in deltas.iter().enumerate() { | |
| self.quants[i].delta_ort_ultimo = *d; | |
| } | |
| Ok((total, PerdaFaseA { | |
| rec: ce.to_scalar::<f32>()? as f64, | |
| commit, | |
| nivel_parada, | |
| acc_recon: acc, | |
| deltas_ort: deltas, | |
| })) | |
| } | |
| /// Decodificação LIVRE (amostragem): ẑ (b, n_ℓ, d_lat) → ids de token AR | |
| /// greedy. Teo A.2/B.4(b): qualquer prefixo decodifica (anytime). | |
| /// O texto final é reconstruído pelo BPE no lado Python. | |
| pub fn decodificar_livre(&self, zhat: &Tensor, max_tokens: usize) -> Result<Vec<Vec<u32>>> { | |
| { | |
| let dev = zhat.device(); | |
| let b = zhat.dim(0)?; | |
| let n_total = self.mascara_causal_dec.dim(0)?; | |
| let up = interp_para(zhat, n_total, dev)?; | |
| let cond = self.dec_proj.fwd(&up)?; | |
| let mut tokens: Vec<Vec<u32>> = (0..b).map(|_| vec![crate::tokens::BOS]).collect(); | |
| for passo in 1..max_tokens.min(n_total) { | |
| let planos: Vec<u32> = tokens.iter().flat_map(|t| t.iter().copied()).collect(); | |
| let entrada = Tensor::from_vec(planos, (b, passo), dev)?; | |
| let mut x = self.dec_emb.fwd(&entrada)?; | |
| x = (&x + &cond.narrow(1, 0, passo)?)?; | |
| for bloco in &self.dec_blocos { | |
| let mut cache = Default::default(); | |
| let m = self.mascara_causal_dec.narrow(0, 0, passo)?.narrow(1, 0, passo)?.contiguous()?; | |
| x = bloco.fwd(&x, Some(&m), &mut cache, None, None)?; | |
| } | |
| let x = self.dec_norm.fwd(&x)?; | |
| let logits = self.dec_head.fwd(&x)?; | |
| let ultimo = logits.narrow(1, passo - 1, 1)?.reshape((b, self.vocab))?; | |
| let preds = ultimo.argmax(D::Minus1)?; | |
| let v = preds.flatten_all()?.to_vec1::<u32>()?; | |
| for (i, t) in v.iter().enumerate() { | |
| tokens[i].push(*t); | |
| } | |
| } | |
| Ok(tokens) | |
| } | |
| } | |
| /// CE por prefixo (V3: monotonia do risco — Teo A.2) em modo congelado. | |
| pub fn ce_por_prefixo(&self, ids: &Tensor) -> Result<Vec<f64>> { | |
| { | |
| let cod = self.codificar_congelado(ids)?; | |
| let mut saida = Vec::with_capacity(cod.zhats.len()); | |
| for zhat in &cod.zhats { | |
| let logits = self.decoder(zhat, Some(ids))?; | |
| let (b, n_t, _) = (logits.dim(0)?, logits.dim(1)?, logits.dim(2)?); | |
| let plano = logits.reshape((b * n_t, self.vocab))?; | |
| let alvos = ids.reshape(b * n_t)?; | |
| let ce = candle_nn::loss::cross_entropy(&plano, &alvos)?; | |
| saida.push(ce.to_scalar::<f32>()? as f64); | |
| } | |
| Ok(saida) | |
| } | |
| } | |
| /// Agenda τ do GSQ por passo (Teo 30.6). | |
| pub fn agendar_tau(&mut self, passo: usize, tau0: f64, tau_min: f64, t_alvo: usize) { | |
| for q in &mut self.quants { | |
| q.agendar_tau(passo, tau0, tau_min, t_alvo); | |
| } | |
| } | |
| /// Atualização EMA do codebook (modo ste_ema apenas — no gsq_rco o | |
| /// gradiente riemanniano já treina W). | |
| pub fn atualizar_ema_codebook(&mut self, ids_lote: &[Tensor]) -> Result<()> { | |
| if self.modo != ModoQuant::SteEma { | |
| return Ok(()); | |
| } | |
| // no modo STE, os codebooks seguem por EMA dos agrupamentos (A.0). | |
| // Implementação: média dos vetores por código → polar → w. | |
| for (i, q) in self.quants.iter_mut().enumerate() { | |
| let ids_l = &ids_lote[i]; | |
| let k = q.uso.len(); | |
| let b = ids_l.dim(0)?; | |
| let plano = ids_l.reshape(b * self.niveis[i])?; | |
| let ids_v = plano.flatten_all()?.to_vec1::<u32>()?; | |
| let mut soma = vec![vec![0f32; self.d_lat]; k]; | |
| let mut conta = vec![0usize; k]; | |
| // usa o codebook corrente para os vetores (aproximação EMA dos | |
| // agrupamentos — suficiente para o modo STE de compatibilidade) | |
| let c = q.codebook()?.detach(); | |
| let c_v = c.flatten_all()?.to_vec1::<f32>()?; | |
| for &idc in &ids_v { | |
| let idc = idc as usize; | |
| if idc < k { | |
| for t in 0..self.d_lat { | |
| soma[idc][t] += c_v[idc * self.d_lat + t]; | |
| } | |
| conta[idc] += 1; | |
| } | |
| } | |
| let ema = q.ema; | |
| let w_v = q.w.as_tensor().flatten_all()?.to_vec1::<f32>()?; | |
| let mut novo = w_v.clone(); | |
| for idc in 0..k { | |
| if conta[idc] > 0 { | |
| for t in 0..self.d_lat { | |
| let alvo = soma[idc][t] / conta[idc] as f32; | |
| novo[idc * self.d_lat + t] = | |
| (ema as f32 * w_v[idc * self.d_lat + t] + (1.0 - ema as f32) * alvo) as f32; | |
| } | |
| } | |
| } | |
| let dev = q.w.as_tensor().device(); | |
| let nt = Tensor::from_vec(novo, (k, self.d_lat), dev)?; | |
| q.w.set(&nt)?; | |
| } | |
| Ok(()) | |
| } | |
| pub fn vars(&self) -> Vec<Var> { | |
| let mut v: Vec<Var> = self.enc_emb.vars(); | |
| for b in &self.enc_blocos { v.extend(b.vars()); } | |
| v.extend(self.enc_norm.vars()); | |
| for p in &self.proj_niveis { v.extend(p.vars()); } | |
| for q in &self.quants { v.push(q.w.clone()); } | |
| v.extend(self.dec_proj.vars()); | |
| v.extend(self.dec_emb.vars()); | |
| for b in &self.dec_blocos { v.extend(b.vars()); } | |
| v.extend(self.dec_norm.vars()); | |
| v.extend(self.dec_head.vars()); | |
| v | |
| } | |
| /// Salva estado (checkpoint atômico — Teo 35.3): índices = ordem de vars(). | |
| pub fn salvar(&self, caminho: &str) -> Result<()> { | |
| let mut obj = serde_json::Map::new(); | |
| for (i, v) in self.vars().iter().enumerate() { | |
| let dados = v.as_tensor().flatten_all()?.to_vec1::<f32>()?; | |
| obj.insert(format!("v{i}"), serde_json::json!({ | |
| "forma": v.as_tensor().dims(), | |
| "dados": dados, | |
| })); | |
| } | |
| let txt = serde_json::to_string(&serde_json::Value::Object(obj)) | |
| .map_err(|e| candle_core::Error::Msg(format!("serde: {e}")))?; | |
| let tmp = format!("{caminho}.tmp"); | |
| std::fs::write(&tmp, txt).map_err(|e| candle_core::Error::Msg(format!("io: {e}")))?; | |
| std::fs::rename(&tmp, caminho).map_err(|e| candle_core::Error::Msg(format!("rename: {e}")))?; | |
| Ok(()) | |
| } | |
| } | |
| pub struct PerdaFaseA { | |
| pub rec: f64, | |
| pub commit: f64, | |
| pub nivel_parada: usize, | |
| pub acc_recon: f64, | |
| pub deltas_ort: Vec<f64>, | |
| } | |
| /// Máscara causal aditiva (n, n): 0 permitido, MIN bloqueado. | |
| pub fn causal_aditiva(n: usize, dev: &Device) -> Result<Tensor> { | |
| let mut m = vec![f32::MIN; n * n]; | |
| for i in 0..n { | |
| for j in 0..=i { | |
| m[i * n + j] = 0.0; | |
| } | |
| } | |
| Tensor::from_vec(m, (n, n), dev) | |
| } | |
| fn tensor_para_json(t: &Tensor) -> Result<serde_json::Value> { | |
| let dados = t.flatten_all()?.to_vec1::<f32>()?; | |
| Ok(serde_json::json!({ | |
| "forma": t.shape().dims(), | |
| "dados": dados, | |
| })) | |
| } | |
| /// Matriz de interpolação cacheada (evita recriar por passo). | |
| pub struct CacheInterp { | |
| pub n_de: usize, | |
| pub n_para: usize, | |
| pub matriz: Tensor, | |
| } | |
| impl CacheInterp { | |
| pub fn novo(n_de: usize, n_para: usize, dev: &Device) -> Result<Self> { | |
| Ok(CacheInterp { n_de, n_para, matriz: matriz_interp(n_de, n_para, dev)? }) | |
| } | |
| } | |