KHTST_RUST v0.1.11 — refatoração PyO3 do KHTST com NCM (provas integradas), ADBA e ponte duplex; treino real Wikipédia PT-BR; suíte 7/7

97df2ae verified 23 days ago

Raw History Blame

17.4 kB

//! KHTST_RUST — refatoração Rust/PyO3 do KHTST com NCM (Next-Concept
//! Modeling), provas integradas (memoria-ncm-khtst n-02/n-03), atenção ADBA
//! (r-04) e ponte duplex (r-05). Documentação: docs/memorias/.
//!
//! Convenções do projeto (n-00 §4): nascimento neutro, leitura honesta,
//! Teos numerados com ∎, token EXCLUSIVAMENTE em HF_TOKEN apagado após uso.

use pyo3::prelude::*;

pub mod amostrador;
pub mod atencao;
pub mod camadas;
pub mod config;
pub mod controlador;
pub mod diagnostico;
pub mod gerador;
pub mod mascaras;
pub mod memoria;
pub mod perdas;
pub mod ponte;
pub mod quantizacao;
pub mod quantizador;
pub mod ssd;
pub mod stiefel;
pub mod telemetria;
pub mod tokens;
pub mod tokenizador;
pub mod treino;
pub mod tronco;

pub mod build_info {
    include!(concat!(env!("OUT_DIR"), "/built.rs"));
}

#[pyfunction]
fn versao() -> &'static str {
    "0.1.11 (KHTST_RUST — NCM + provas integradas + ADBA + ponte duplex)"
}

#[pyfunction]
fn info_build() -> String {
    format!(
        "rustc={}; perfil={}; features={}; alvo={}",
        build_info::RUSTC_VERSION,
        build_info::PROFILE,
        build_info::FEATURES_STR,
        build_info::TARGET,
    )
}

/// Lint de config com sugestões por Levenshtein (diagnóstico de digitação).
#[pyfunction]
fn lint_config(py: Python, caminho: &str) -> PyResult<PyObject> {
    let txt = std::fs::read_to_string(caminho)
        .map_err(|e| pyo3::exceptions::PyIOError::new_err(format!("{caminho}: {e}")))?;
    let v: serde_json::Value = serde_json::from_str(&txt)
        .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("JSON inválido: {e}")))?;
    let achatado = diagnostico::achatar(&v);
    let avisos = diagnostico::lint_chaves(&achatado);
    let saida: Vec<serde_json::Value> = avisos
        .iter()
        .map(|a| serde_json::json!({
            "chave": a.chave, "sugestao": a.sugestao, "causa": a.causa, "correcao": a.correcao,
        }))
        .collect();
    Ok(pythonize::pythonize(py, &serde_json::json!(saida))?.into())
}

#[pyclass(name = "Bpe")]
struct BpePy {
    bpe: tokens::Bpe,
}

#[pymethods]
impl BpePy {
    #[staticmethod]
    fn treinar(textos: Vec<String>, vocab_alvo: usize) -> PyResult<Self> {
        Ok(BpePy { bpe: tokens::Bpe::treinar(&textos, vocab_alvo) })
    }

    #[staticmethod]
    fn carregar(caminho: &str) -> PyResult<Self> {
        let txt = std::fs::read_to_string(caminho)
            .map_err(|e| pyo3::exceptions::PyIOError::new_err(format!("{e}")))?;
        let v: serde_json::Value = serde_json::from_str(&txt)
            .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("JSON inválido: {e}")))?;
        let bpe = tokens::Bpe::de_json(&v)
            .ok_or_else(|| pyo3::exceptions::PyValueError::new_err("formato BPE inválido"))?;
        Ok(BpePy { bpe })
    }

    fn codificar(&mut self, texto: &str) -> Vec<u32> {
        self.bpe.codificar(texto)
    }

    fn decodificar(&self, ids: Vec<u32>) -> String {
        self.bpe.decodificar(&ids)
    }

    fn salvar(&self, caminho: &str) -> PyResult<()> {
        std::fs::write(caminho, serde_json::to_string(&self.bpe.para_json()).unwrap())
            .map_err(|e| pyo3::exceptions::PyIOError::new_err(format!("{e}")))?;
        Ok(())
    }

    #[getter]
    fn n_vocab(&self) -> usize {
        self.bpe.n_vocab()
    }
}

#[pyclass(name = "Config")]
struct ConfigPy {
    cfg: config::ConfigKhtst,
    caminho: Option<String>,
}

#[pymethods]
impl ConfigPy {
    #[new]
    #[allow(never_type_fallback)]
    #[pyo3(signature = (caminho=None))]
    fn new(caminho: Option<String>) -> PyResult<Self> {
        let cfg = match &caminho {
            Some(c) => {
                let txt = std::fs::read_to_string(c)
                    .map_err(|e| pyo3::exceptions::PyIOError::new_err(format!("{c}: {e}")))?;
                let v: serde_json::Value = serde_json::from_str(&txt)
                    .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("JSON inválido: {e}")))?;
                serde_json::from_value(v)
                    .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("config: {e}")))?
            }
            None => config::ConfigKhtst::default(),
        };
        cfg.validar().map_err(pyo3::exceptions::PyValueError::new_err)?;
        Ok(ConfigPy { cfg, caminho })
    }

    fn validar(&self) -> PyResult<()> {
        self.cfg.validar().map_err(pyo3::exceptions::PyValueError::new_err)
    }

    fn checar_orcamento_cache(&self) -> PyResult<usize> {
        self.cfg.checar_orcamento_cache().map_err(pyo3::exceptions::PyValueError::new_err)
    }

    fn resumo(&self) -> String {
        format!(
            "modelo(vocab={}, d={}, camadas={}, ctx={}) ncm(L={}, niveis={:?}, K={}, d_lat={}, modo={}, ativo={})",
            self.cfg.modelo.vocab, self.cfg.modelo.d_modelo, self.cfg.modelo.n_camadas,
            self.cfg.modelo.comprimento_ctx, self.cfg.l_niveis(), self.cfg.ncm.niveis,
            self.cfg.ncm.k_por_nivel, self.cfg.ncm.d_lat, self.cfg.ncm.modo_quantizacao, self.cfg.ncm.ativo,
        )
    }

    #[getter]
    fn l_niveis(&self) -> usize {
        self.cfg.l_niveis()
    }
}

#[pyclass(name = "PonteDupla")]
struct PonteDuplaPy {
    ponte: ponte::PonteDupla,
}

#[pymethods]
impl PonteDuplaPy {
    #[new]
    #[pyo3(signature = (capacidade=64, timeout_ms=2000))]
    fn new(capacidade: usize, timeout_ms: u64) -> Self {
        PonteDuplaPy { ponte: ponte::PonteDupla::nova(capacidade, timeout_ms) }
    }

    fn rodada(&self, carga_a: Vec<f32>, carga_b: Vec<f32>, nivel: usize) -> PyResult<(Vec<f32>, Vec<f32>)> {
        let (ma, mb) = self.ponte.rodada(carga_a, carga_b, nivel)
            .map_err(pyo3::exceptions::PyRuntimeError::new_err)?;
        Ok((ma.carga, mb.carga))
    }

    fn demo_professor_aluno(&self, rodadas: usize, dim: usize) -> PyResult<PyObject> {
        let inicio: Vec<f32> = (0..dim).map(|i| (i as f32) * 0.01).collect();
        let stats = ponte::demo_professor_aluno(&self.ponte, rodadas, inicio, &|a, _| {
            a.iter().map(|x| x * 1.05).collect()
        })
        .map_err(pyo3::exceptions::PyRuntimeError::new_err)?;
        Python::with_gil(|py| {
            Ok(pythonize::pythonize(py, &serde_json::to_value(&stats).unwrap())?.into())
        })
    }

    fn estatisticas(&self) -> PyResult<PyObject> {
        let s = self.ponte.estatisticas();
        Python::with_gil(|py| {
            Ok(pythonize::pythonize(py, &serde_json::to_value(&s).unwrap())?.into())
        })
    }
}

#[cfg(feature = "ponte_async")]
#[pyclass(name = "PonteAsync")]
struct PonteAsyncPy {
    interna: std::sync::Arc<ponte::assincrona::PonteAsync>,
}

#[cfg(feature = "ponte_async")]
#[pymethods]
impl PonteAsyncPy {
    #[new]
    #[pyo3(signature = (capacidade=64, timeout_ms=2000))]
    fn new(capacidade: usize, timeout_ms: u64) -> Self {
        PonteAsyncPy { interna: std::sync::Arc::new(ponte::assincrona::PonteAsync::nova(capacidade, timeout_ms)) }
    }

    /// Corrotina: `await ponte.trocar(a, b, nivel)` — duplex simultânea.
    fn trocar<'p>(&self, py: Python<'p>, carga_a: Vec<f32>, carga_b: Vec<f32>, nivel: usize) -> PyResult<Bound<'p, PyAny>> {
        let p = self.interna.clone();
        pyo3_async_runtimes::tokio::future_into_py(py, async move {
            let (ma, mb) = p.rodada_async(carga_a, carga_b, nivel)
                .await
                .map_err(pyo3::exceptions::PyRuntimeError::new_err)?;
            Ok::<_, PyErr>((ma.carga, mb.carga))
        })
    }
}

/// Treina a fase A (tokenizador). Devolve dict com resumo + telemetria.
#[pyfunction]
#[pyo3(signature = (cfg, sequencias, caminho_telemetria=None, caminho_checkpoint=None, callback=None, teto_seg=None))]
fn treinar_fase_a(
    cfg: &ConfigPy,
    sequencias: Vec<Vec<u32>>,
    caminho_telemetria: Option<String>,
    caminho_checkpoint: Option<String>,
    callback: Option<Bound<'_, PyAny>>,
    teto_seg: Option<f64>,
) -> PyResult<PyObject> {
    let cb: treino::CallbackPy = callback.map(|c| c.unbind());
    let mut hub = telemetria::TelemetryHub::novo(caminho_telemetria);
    let dados = treino::DadosTreino::novo(sequencias, cfg.cfg.modelo.comprimento_ctx);
    let mut resumo = treino::treinar_fase_a(&cfg.cfg, &dados, &mut hub, &cb, caminho_checkpoint.as_deref(), teto_seg)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("fase A: {e}")))?;
    let _ = &mut resumo;
    Python::with_gil(|py| {
        let dados = serde_json::json!({
            "resumo": {
                "passos": resumo.passos, "perda_inicial": resumo.perda_inicial,
                "perda_final": resumo.perda_final, "acc_recon": resumo.acc_recon,
                "delta_ort_max": resumo.delta_ort_max,
                "perplexidade_codebook": resumo.perplexidade_codebook,
                "monotonia_prefixo": resumo.monotonia_prefixo,
                "escore_aprendizado": resumo.escore_aprendizado,
            },
            "telemetria": hub.dump(),
        });
        Ok(pythonize::pythonize(py, &dados)?.into())
    })
}

/// Treina a fase B (gerador) com tokenizador congelado carregado de checkpoint.
#[pyfunction]
#[pyo3(signature = (cfg, sequencias, caminho_tokenizador, caminho_telemetria=None, caminho_checkpoint=None, callback=None, teto_seg=None))]
fn treinar_fase_b(
    cfg: &ConfigPy,
    sequencias: Vec<Vec<u32>>,
    caminho_tokenizador: &str,
    caminho_telemetria: Option<String>,
    caminho_checkpoint: Option<String>,
    callback: Option<Bound<'_, PyAny>>,
    teto_seg: Option<f64>,
) -> PyResult<PyObject> {
    let cb: treino::CallbackPy = callback.map(|c| c.unbind());
    let mut hub = telemetria::TelemetryHub::novo(caminho_telemetria);
    let dados = treino::DadosTreino::novo(sequencias, cfg.cfg.modelo.comprimento_ctx);
    // carrega o tokenizador congelado
    let dev = candle_core::Device::Cpu;
    let mut tok = tokenizador::TokenizadorConceitos::novo(&cfg.cfg, &dev)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("construção: {e}")))?;
    carregar_tokenizador(&mut tok, caminho_tokenizador)
        .map_err(pyo3::exceptions::PyRuntimeError::new_err)?;
    let resumo = treino::treinar_fase_b(&cfg.cfg, &tok, &dados, &mut hub, &cb, caminho_checkpoint.as_deref(), teto_seg)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("fase B: {e}")))?;
    Python::with_gil(|py| {
        let dados = serde_json::json!({
            "resumo": {
                "passos": resumo.passos, "perda_inicial": resumo.perda_inicial,
                "perda_final": resumo.perda_final, "acc_conceitos": resumo.acc_recon,
                "escore_aprendizado": resumo.escore_aprendizado,
            },
            "telemetria": hub.dump(),
        });
        Ok(pythonize::pythonize(py, &dados)?.into())
    })
}

fn carregar_tokenizador(tok: &mut tokenizador::TokenizadorConceitos, caminho: &str) -> Result<(), String> {
    let txt = std::fs::read_to_string(caminho).map_err(|e| format!("{e}"))?;
    let v: serde_json::Value = serde_json::from_str(&txt).map_err(|e| format!("{e}"))?;
    let obj = v.as_object().ok_or("formato de checkpoint inválido")?;
    let vars = tok.vars();
    for (i, var) in vars.iter().enumerate() {
        if let Some(entrada) = obj.get(&format!("v{i}")) {
            let forma: Vec<usize> = entrada["forma"].as_array().map(|a| a.iter().filter_map(|x| x.as_u64().map(|y| y as usize)).collect()).unwrap_or_default();
            let dados: Vec<f32> = entrada["dados"].as_array().map(|a| a.iter().filter_map(|x| x.as_f64().map(|y| y as f32)).collect()).unwrap_or_default();
            if forma.is_empty() || dados.is_empty() {
                continue;
            }
            let dev = candle_core::Device::Cpu;
            let t2 = candle_core::Tensor::from_vec(dados, forma.clone(), &dev)
                .and_then(|t| t.reshape(forma))
                .map_err(|e| format!("{e}"))?;
            var.set(&t2).map_err(|e| format!("{e}"))?;
        }
    }
    Ok(())
}

/// Gera amostras com NFE = L e devolve textos decodificados (modo conceitos).
#[pyfunction]
#[pyo3(signature = (cfg, caminho_tokenizador, caminho_gerador, n_amostras=2, parar_no_nivel=None, greedy=true))]
fn gerar_amostras(
    cfg: &ConfigPy,
    caminho_tokenizador: &str,
    caminho_gerador: &str,
    n_amostras: usize,
    parar_no_nivel: Option<usize>,
    greedy: bool,
) -> PyResult<PyObject> {
    let dev = candle_core::Device::Cpu;
    let mut tok = tokenizador::TokenizadorConceitos::novo(&cfg.cfg, &dev)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    carregar_tokenizador(&mut tok, caminho_tokenizador)
        .map_err(pyo3::exceptions::PyRuntimeError::new_err)?;
    let mut ger = gerador::GeradorConceitos::novo(&cfg.cfg, &dev)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    carregar_gerador(&mut ger, caminho_gerador)
        .map_err(pyo3::exceptions::PyRuntimeError::new_err)?;
    let mut amostrador = amostrador::AmostradorConceitos::novo(
        &mut ger, &tok, cfg.cfg.ncm.cfg_scale, cfg.cfg.ncm.temperaturas.limites, cfg.cfg.ncm.temperaturas.entropia_alvo,
    );
    let mut rng = <rand::prelude::StdRng as rand::SeedableRng>::seed_from_u64(4242);
    let inicio = std::time::Instant::now();
    let (codigos, nfe) = amostrador.gerar(n_amostras.max(1), &dev, &mut rng, parar_no_nivel, greedy)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("amostragem: {e}")))?;
    let tempo_ms = inicio.elapsed().as_secs_f64() * 1000.0;
    // decodifica: ẑ final → decoder AR greedy → tokens
    let nivel_final = codigos.len();
    let zhats = tok.reconstruir_de_codigos(&codigos, nivel_final, n_amostras.max(1), &dev)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    let textos_ids = decodificar_textos(&tok, zhats.last(), n_amostras.max(1))
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    Python::with_gil(|py| {
        let d = serde_json::json!({
            "nfe": nfe, "l": cfg.cfg.l_niveis(), "tempo_ms": tempo_ms,
            "niveis_usados": nivel_final, "ids_tokens": textos_ids,
            "temperaturas": amostrador.temperaturas,
        });
        Ok(pythonize::pythonize(py, &d)?.into())
    })
}

fn decodificar_textos(tok: &tokenizador::TokenizadorConceitos, zhat: Option<&candle_core::Tensor>, b: usize) -> Result<Vec<Vec<u32>>, String> {
    let Some(z) = zhat else { return Ok(vec![vec![]; b]) };
    tok.decodificar_livre(z, 48).map_err(|e| format!("{e}"))
}

fn carregar_gerador(ger: &mut gerador::GeradorConceitos, caminho: &str) -> Result<(), String> {
    let txt = std::fs::read_to_string(caminho).map_err(|e| format!("{e}"))?;
    let v: serde_json::Value = serde_json::from_str(&txt).map_err(|e| format!("{e}"))?;
    let obj = v.as_object().ok_or("formato inválido")?;
    let vars = ger.vars();
    for (i, var) in vars.iter().enumerate() {
        if let Some(entrada) = obj.get(&format!("v{i}")) {
            let forma: Vec<usize> = entrada["forma"].as_array().map(|a| a.iter().filter_map(|x| x.as_u64().map(|y| y as usize)).collect()).unwrap_or_default();
            let dados: Vec<f32> = entrada["dados"].as_array().map(|a| a.iter().filter_map(|x| x.as_f64().map(|y| y as f32)).collect()).unwrap_or_default();
            if forma.is_empty() || dados.is_empty() {
                continue;
            }
            let dev = candle_core::Device::Cpu;
            let t = candle_core::Tensor::from_vec(dados, forma.clone(), &dev)
                .and_then(|t| t.reshape(forma))
                .map_err(|e| format!("{e}"))?;
            var.set(&t).map_err(|e| format!("{e}"))?;
        }
    }
    Ok(())
}

/// Utilitários de prova expostos ao Python (V-série da suíte).
#[pyfunction]
fn verificar_mascara(cfg: &ConfigPy) -> PyResult<usize> {
    let dev = candle_core::Device::Cpu;
    let m = mascaras::mascara_block_causal(&cfg.cfg.ncm.niveis, gerador::N_PREFIXO, &dev)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    mascaras::contar_vazamentos(&cfg.cfg.ncm.niveis, gerador::N_PREFIXO, &m)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))
}

#[pyfunction]
fn verificar_stiefel(k: usize, d: usize) -> PyResult<f64> {
    let dev = candle_core::Device::Cpu;
    let c = stiefel::init_stiefel(k, d, &dev)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    stiefel::erro_ortogonalidade(&c.transpose(0, 1).unwrap()).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))
}

#[pymodule]
fn nucleo_rust(m: &Bound<'_, PyModule>) -> PyResult<()> {
    pyo3_log::init();
    m.add_function(wrap_pyfunction!(versao, m)?)?;
    m.add_function(wrap_pyfunction!(info_build, m)?)?;
    m.add_function(wrap_pyfunction!(lint_config, m)?)?;
    m.add_function(wrap_pyfunction!(treinar_fase_a, m)?)?;
    m.add_function(wrap_pyfunction!(treinar_fase_b, m)?)?;
    m.add_function(wrap_pyfunction!(gerar_amostras, m)?)?;
    m.add_function(wrap_pyfunction!(verificar_mascara, m)?)?;
    m.add_function(wrap_pyfunction!(verificar_stiefel, m)?)?;
    m.add_class::<BpePy>()?;
    m.add_class::<ConfigPy>()?;
    m.add_class::<PonteDuplaPy>()?;
    #[cfg(feature = "ponte_async")]
    m.add_class::<PonteAsyncPy>()?;
    Ok(())
}

	//! KHTST_RUST — refatoração Rust/PyO3 do KHTST com NCM (Next-Concept
	//! Modeling), provas integradas (memoria-ncm-khtst n-02/n-03), atenção ADBA
	//! (r-04) e ponte duplex (r-05). Documentação: docs/memorias/.
	//!
	//! Convenções do projeto (n-00 §4): nascimento neutro, leitura honesta,
	//! Teos numerados com ∎, token EXCLUSIVAMENTE em HF_TOKEN apagado após uso.

	use pyo3::prelude::*;

	pub mod amostrador;
	pub mod atencao;
	pub mod camadas;
	pub mod config;
	pub mod controlador;
	pub mod diagnostico;
	pub mod gerador;
	pub mod mascaras;
	pub mod memoria;
	pub mod perdas;
	pub mod ponte;
	pub mod quantizacao;
	pub mod quantizador;
	pub mod ssd;
	pub mod stiefel;
	pub mod telemetria;
	pub mod tokens;
	pub mod tokenizador;
	pub mod treino;
	pub mod tronco;

	pub mod build_info {
	include!(concat!(env!("OUT_DIR"), "/built.rs"));
	}

	#[pyfunction]
	fn versao() -> &'static str {
	"0.1.11 (KHTST_RUST — NCM + provas integradas + ADBA + ponte duplex)"
	}

	#[pyfunction]
	fn info_build() -> String {
	format!(
	"rustc={}; perfil={}; features={}; alvo={}",
	build_info::RUSTC_VERSION,
	build_info::PROFILE,
	build_info::FEATURES_STR,
	build_info::TARGET,
	)
	}

	/// Lint de config com sugestões por Levenshtein (diagnóstico de digitação).
	#[pyfunction]
	fn lint_config(py: Python, caminho: &str) -> PyResult<PyObject> {
	let txt = std::fs::read_to_string(caminho)
	.map_err(\|e\| pyo3::exceptions::PyIOError::new_err(format!("{caminho}: {e}")))?;
	let v: serde_json::Value = serde_json::from_str(&txt)
	.map_err(\|e\| pyo3::exceptions::PyValueError::new_err(format!("JSON inválido: {e}")))?;
	let achatado = diagnostico::achatar(&v);
	let avisos = diagnostico::lint_chaves(&achatado);
	let saida: Vec<serde_json::Value> = avisos
	.iter()
	.map(\|a\| serde_json::json!({
	"chave": a.chave, "sugestao": a.sugestao, "causa": a.causa, "correcao": a.correcao,
	}))
	.collect();
	Ok(pythonize::pythonize(py, &serde_json::json!(saida))?.into())
	}

	#[pyclass(name = "Bpe")]
	struct BpePy {
	bpe: tokens::Bpe,
	}

	#[pymethods]
	impl BpePy {
	#[staticmethod]
	fn treinar(textos: Vec<String>, vocab_alvo: usize) -> PyResult<Self> {
	Ok(BpePy { bpe: tokens::Bpe::treinar(&textos, vocab_alvo) })
	}

	#[staticmethod]
	fn carregar(caminho: &str) -> PyResult<Self> {
	let txt = std::fs::read_to_string(caminho)
	.map_err(\|e\| pyo3::exceptions::PyIOError::new_err(format!("{e}")))?;
	let v: serde_json::Value = serde_json::from_str(&txt)
	.map_err(\|e\| pyo3::exceptions::PyValueError::new_err(format!("JSON inválido: {e}")))?;
	let bpe = tokens::Bpe::de_json(&v)
	.ok_or_else(\|\| pyo3::exceptions::PyValueError::new_err("formato BPE inválido"))?;
	Ok(BpePy { bpe })
	}

	fn codificar(&mut self, texto: &str) -> Vec<u32> {
	self.bpe.codificar(texto)
	}

	fn decodificar(&self, ids: Vec<u32>) -> String {
	self.bpe.decodificar(&ids)
	}

	fn salvar(&self, caminho: &str) -> PyResult<()> {
	std::fs::write(caminho, serde_json::to_string(&self.bpe.para_json()).unwrap())
	.map_err(\|e\| pyo3::exceptions::PyIOError::new_err(format!("{e}")))?;
	Ok(())
	}

	#[getter]
	fn n_vocab(&self) -> usize {
	self.bpe.n_vocab()
	}
	}

	#[pyclass(name = "Config")]
	struct ConfigPy {
	cfg: config::ConfigKhtst,
	caminho: Option<String>,
	}

	#[pymethods]
	impl ConfigPy {
	#[new]
	#[allow(never_type_fallback)]
	#[pyo3(signature = (caminho=None))]
	fn new(caminho: Option<String>) -> PyResult<Self> {
	let cfg = match &caminho {
	Some(c) => {
	let txt = std::fs::read_to_string(c)
	.map_err(\|e\| pyo3::exceptions::PyIOError::new_err(format!("{c}: {e}")))?;
	let v: serde_json::Value = serde_json::from_str(&txt)
	.map_err(\|e\| pyo3::exceptions::PyValueError::new_err(format!("JSON inválido: {e}")))?;
	serde_json::from_value(v)
	.map_err(\|e\| pyo3::exceptions::PyValueError::new_err(format!("config: {e}")))?
	}
	None => config::ConfigKhtst::default(),
	};
	cfg.validar().map_err(pyo3::exceptions::PyValueError::new_err)?;
	Ok(ConfigPy { cfg, caminho })
	}

	fn validar(&self) -> PyResult<()> {
	self.cfg.validar().map_err(pyo3::exceptions::PyValueError::new_err)
	}

	fn checar_orcamento_cache(&self) -> PyResult<usize> {
	self.cfg.checar_orcamento_cache().map_err(pyo3::exceptions::PyValueError::new_err)
	}

	fn resumo(&self) -> String {
	format!(
	"modelo(vocab={}, d={}, camadas={}, ctx={}) ncm(L={}, niveis={:?}, K={}, d_lat={}, modo={}, ativo={})",
	self.cfg.modelo.vocab, self.cfg.modelo.d_modelo, self.cfg.modelo.n_camadas,
	self.cfg.modelo.comprimento_ctx, self.cfg.l_niveis(), self.cfg.ncm.niveis,
	self.cfg.ncm.k_por_nivel, self.cfg.ncm.d_lat, self.cfg.ncm.modo_quantizacao, self.cfg.ncm.ativo,
	)
	}

	#[getter]
	fn l_niveis(&self) -> usize {
	self.cfg.l_niveis()
	}
	}

	#[pyclass(name = "PonteDupla")]
	struct PonteDuplaPy {
	ponte: ponte::PonteDupla,
	}

	#[pymethods]
	impl PonteDuplaPy {
	#[new]
	#[pyo3(signature = (capacidade=64, timeout_ms=2000))]
	fn new(capacidade: usize, timeout_ms: u64) -> Self {
	PonteDuplaPy { ponte: ponte::PonteDupla::nova(capacidade, timeout_ms) }
	}

	fn rodada(&self, carga_a: Vec<f32>, carga_b: Vec<f32>, nivel: usize) -> PyResult<(Vec<f32>, Vec<f32>)> {
	let (ma, mb) = self.ponte.rodada(carga_a, carga_b, nivel)
	.map_err(pyo3::exceptions::PyRuntimeError::new_err)?;
	Ok((ma.carga, mb.carga))
	}

	fn demo_professor_aluno(&self, rodadas: usize, dim: usize) -> PyResult<PyObject> {
	let inicio: Vec<f32> = (0..dim).map(\|i\| (i as f32) * 0.01).collect();
	let stats = ponte::demo_professor_aluno(&self.ponte, rodadas, inicio, &\|a, _\| {
	a.iter().map(\|x\| x * 1.05).collect()
	})
	.map_err(pyo3::exceptions::PyRuntimeError::new_err)?;
	Python::with_gil(\|py\| {
	Ok(pythonize::pythonize(py, &serde_json::to_value(&stats).unwrap())?.into())
	})
	}

	fn estatisticas(&self) -> PyResult<PyObject> {
	let s = self.ponte.estatisticas();
	Python::with_gil(\|py\| {
	Ok(pythonize::pythonize(py, &serde_json::to_value(&s).unwrap())?.into())
	})
	}
	}

	#[cfg(feature = "ponte_async")]
	#[pyclass(name = "PonteAsync")]
	struct PonteAsyncPy {
	interna: std::sync::Arc<ponte::assincrona::PonteAsync>,
	}

	#[cfg(feature = "ponte_async")]
	#[pymethods]
	impl PonteAsyncPy {
	#[new]
	#[pyo3(signature = (capacidade=64, timeout_ms=2000))]
	fn new(capacidade: usize, timeout_ms: u64) -> Self {
	PonteAsyncPy { interna: std::sync::Arc::new(ponte::assincrona::PonteAsync::nova(capacidade, timeout_ms)) }
	}

	/// Corrotina: `await ponte.trocar(a, b, nivel)` — duplex simultânea.
	fn trocar<'p>(&self, py: Python<'p>, carga_a: Vec<f32>, carga_b: Vec<f32>, nivel: usize) -> PyResult<Bound<'p, PyAny>> {
	let p = self.interna.clone();
	pyo3_async_runtimes::tokio::future_into_py(py, async move {
	let (ma, mb) = p.rodada_async(carga_a, carga_b, nivel)
	.await
	.map_err(pyo3::exceptions::PyRuntimeError::new_err)?;
	Ok::<_, PyErr>((ma.carga, mb.carga))
	})
	}
	}

	/// Treina a fase A (tokenizador). Devolve dict com resumo + telemetria.
	#[pyfunction]
	#[pyo3(signature = (cfg, sequencias, caminho_telemetria=None, caminho_checkpoint=None, callback=None, teto_seg=None))]
	fn treinar_fase_a(
	cfg: &ConfigPy,
	sequencias: Vec<Vec<u32>>,
	caminho_telemetria: Option<String>,
	caminho_checkpoint: Option<String>,
	callback: Option<Bound<'_, PyAny>>,
	teto_seg: Option<f64>,
	) -> PyResult<PyObject> {
	let cb: treino::CallbackPy = callback.map(\|c\| c.unbind());
	let mut hub = telemetria::TelemetryHub::novo(caminho_telemetria);
	let dados = treino::DadosTreino::novo(sequencias, cfg.cfg.modelo.comprimento_ctx);
	let mut resumo = treino::treinar_fase_a(&cfg.cfg, &dados, &mut hub, &cb, caminho_checkpoint.as_deref(), teto_seg)
	.map_err(\|e\| pyo3::exceptions::PyRuntimeError::new_err(format!("fase A: {e}")))?;
	let _ = &mut resumo;
	Python::with_gil(\|py\| {
	let dados = serde_json::json!({
	"resumo": {
	"passos": resumo.passos, "perda_inicial": resumo.perda_inicial,
	"perda_final": resumo.perda_final, "acc_recon": resumo.acc_recon,
	"delta_ort_max": resumo.delta_ort_max,
	"perplexidade_codebook": resumo.perplexidade_codebook,
	"monotonia_prefixo": resumo.monotonia_prefixo,
	"escore_aprendizado": resumo.escore_aprendizado,
	},
	"telemetria": hub.dump(),
	});
	Ok(pythonize::pythonize(py, &dados)?.into())
	})
	}

	/// Treina a fase B (gerador) com tokenizador congelado carregado de checkpoint.
	#[pyfunction]
	#[pyo3(signature = (cfg, sequencias, caminho_tokenizador, caminho_telemetria=None, caminho_checkpoint=None, callback=None, teto_seg=None))]
	fn treinar_fase_b(
	cfg: &ConfigPy,
	sequencias: Vec<Vec<u32>>,
	caminho_tokenizador: &str,
	caminho_telemetria: Option<String>,
	caminho_checkpoint: Option<String>,
	callback: Option<Bound<'_, PyAny>>,
	teto_seg: Option<f64>,
	) -> PyResult<PyObject> {
	let cb: treino::CallbackPy = callback.map(\|c\| c.unbind());
	let mut hub = telemetria::TelemetryHub::novo(caminho_telemetria);
	let dados = treino::DadosTreino::novo(sequencias, cfg.cfg.modelo.comprimento_ctx);
	// carrega o tokenizador congelado
	let dev = candle_core::Device::Cpu;
	let mut tok = tokenizador::TokenizadorConceitos::novo(&cfg.cfg, &dev)
	.map_err(\|e\| pyo3::exceptions::PyRuntimeError::new_err(format!("construção: {e}")))?;
	carregar_tokenizador(&mut tok, caminho_tokenizador)
	.map_err(pyo3::exceptions::PyRuntimeError::new_err)?;
	let resumo = treino::treinar_fase_b(&cfg.cfg, &tok, &dados, &mut hub, &cb, caminho_checkpoint.as_deref(), teto_seg)
	.map_err(\|e\| pyo3::exceptions::PyRuntimeError::new_err(format!("fase B: {e}")))?;
	Python::with_gil(\|py\| {
	let dados = serde_json::json!({
	"resumo": {
	"passos": resumo.passos, "perda_inicial": resumo.perda_inicial,
	"perda_final": resumo.perda_final, "acc_conceitos": resumo.acc_recon,
	"escore_aprendizado": resumo.escore_aprendizado,
	},
	"telemetria": hub.dump(),
	});
	Ok(pythonize::pythonize(py, &dados)?.into())
	})
	}

	fn carregar_tokenizador(tok: &mut tokenizador::TokenizadorConceitos, caminho: &str) -> Result<(), String> {
	let txt = std::fs::read_to_string(caminho).map_err(\|e\| format!("{e}"))?;
	let v: serde_json::Value = serde_json::from_str(&txt).map_err(\|e\| format!("{e}"))?;
	let obj = v.as_object().ok_or("formato de checkpoint inválido")?;
	let vars = tok.vars();
	for (i, var) in vars.iter().enumerate() {
	if let Some(entrada) = obj.get(&format!("v{i}")) {
	let forma: Vec<usize> = entrada["forma"].as_array().map(\|a\| a.iter().filter_map(\|x\| x.as_u64().map(\|y\| y as usize)).collect()).unwrap_or_default();
	let dados: Vec<f32> = entrada["dados"].as_array().map(\|a\| a.iter().filter_map(\|x\| x.as_f64().map(\|y\| y as f32)).collect()).unwrap_or_default();
	if forma.is_empty() \|\| dados.is_empty() {
	continue;
	}
	let dev = candle_core::Device::Cpu;
	let t2 = candle_core::Tensor::from_vec(dados, forma.clone(), &dev)
	.and_then(\|t\| t.reshape(forma))
	.map_err(\|e\| format!("{e}"))?;
	var.set(&t2).map_err(\|e\| format!("{e}"))?;
	}
	}
	Ok(())
	}

	/// Gera amostras com NFE = L e devolve textos decodificados (modo conceitos).
	#[pyfunction]
	#[pyo3(signature = (cfg, caminho_tokenizador, caminho_gerador, n_amostras=2, parar_no_nivel=None, greedy=true))]
	fn gerar_amostras(
	cfg: &ConfigPy,
	caminho_tokenizador: &str,
	caminho_gerador: &str,
	n_amostras: usize,
	parar_no_nivel: Option<usize>,
	greedy: bool,
	) -> PyResult<PyObject> {
	let dev = candle_core::Device::Cpu;
	let mut tok = tokenizador::TokenizadorConceitos::novo(&cfg.cfg, &dev)
	.map_err(\|e\| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
	carregar_tokenizador(&mut tok, caminho_tokenizador)
	.map_err(pyo3::exceptions::PyRuntimeError::new_err)?;
	let mut ger = gerador::GeradorConceitos::novo(&cfg.cfg, &dev)
	.map_err(\|e\| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
	carregar_gerador(&mut ger, caminho_gerador)
	.map_err(pyo3::exceptions::PyRuntimeError::new_err)?;
	let mut amostrador = amostrador::AmostradorConceitos::novo(
	&mut ger, &tok, cfg.cfg.ncm.cfg_scale, cfg.cfg.ncm.temperaturas.limites, cfg.cfg.ncm.temperaturas.entropia_alvo,
	);
	let mut rng = <rand::prelude::StdRng as rand::SeedableRng>::seed_from_u64(4242);
	let inicio = std::time::Instant::now();
	let (codigos, nfe) = amostrador.gerar(n_amostras.max(1), &dev, &mut rng, parar_no_nivel, greedy)
	.map_err(\|e\| pyo3::exceptions::PyRuntimeError::new_err(format!("amostragem: {e}")))?;
	let tempo_ms = inicio.elapsed().as_secs_f64() * 1000.0;
	// decodifica: ẑ final → decoder AR greedy → tokens
	let nivel_final = codigos.len();
	let zhats = tok.reconstruir_de_codigos(&codigos, nivel_final, n_amostras.max(1), &dev)
	.map_err(\|e\| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
	let textos_ids = decodificar_textos(&tok, zhats.last(), n_amostras.max(1))
	.map_err(\|e\| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
	Python::with_gil(\|py\| {
	let d = serde_json::json!({
	"nfe": nfe, "l": cfg.cfg.l_niveis(), "tempo_ms": tempo_ms,
	"niveis_usados": nivel_final, "ids_tokens": textos_ids,
	"temperaturas": amostrador.temperaturas,
	});
	Ok(pythonize::pythonize(py, &d)?.into())
	})
	}

	fn decodificar_textos(tok: &tokenizador::TokenizadorConceitos, zhat: Option<&candle_core::Tensor>, b: usize) -> Result<Vec<Vec<u32>>, String> {
	let Some(z) = zhat else { return Ok(vec![vec![]; b]) };
	tok.decodificar_livre(z, 48).map_err(\|e\| format!("{e}"))
	}

	fn carregar_gerador(ger: &mut gerador::GeradorConceitos, caminho: &str) -> Result<(), String> {
	let txt = std::fs::read_to_string(caminho).map_err(\|e\| format!("{e}"))?;
	let v: serde_json::Value = serde_json::from_str(&txt).map_err(\|e\| format!("{e}"))?;
	let obj = v.as_object().ok_or("formato inválido")?;
	let vars = ger.vars();
	for (i, var) in vars.iter().enumerate() {
	if let Some(entrada) = obj.get(&format!("v{i}")) {
	let forma: Vec<usize> = entrada["forma"].as_array().map(\|a\| a.iter().filter_map(\|x\| x.as_u64().map(\|y\| y as usize)).collect()).unwrap_or_default();
	let dados: Vec<f32> = entrada["dados"].as_array().map(\|a\| a.iter().filter_map(\|x\| x.as_f64().map(\|y\| y as f32)).collect()).unwrap_or_default();
	if forma.is_empty() \|\| dados.is_empty() {
	continue;
	}
	let dev = candle_core::Device::Cpu;
	let t = candle_core::Tensor::from_vec(dados, forma.clone(), &dev)
	.and_then(\|t\| t.reshape(forma))
	.map_err(\|e\| format!("{e}"))?;
	var.set(&t).map_err(\|e\| format!("{e}"))?;
	}
	}
	Ok(())
	}

	/// Utilitários de prova expostos ao Python (V-série da suíte).
	#[pyfunction]
	fn verificar_mascara(cfg: &ConfigPy) -> PyResult<usize> {
	let dev = candle_core::Device::Cpu;
	let m = mascaras::mascara_block_causal(&cfg.cfg.ncm.niveis, gerador::N_PREFIXO, &dev)
	.map_err(\|e\| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
	mascaras::contar_vazamentos(&cfg.cfg.ncm.niveis, gerador::N_PREFIXO, &m)
	.map_err(\|e\| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))
	}

	#[pyfunction]
	fn verificar_stiefel(k: usize, d: usize) -> PyResult<f64> {
	let dev = candle_core::Device::Cpu;
	let c = stiefel::init_stiefel(k, d, &dev)
	.map_err(\|e\| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
	stiefel::erro_ortogonalidade(&c.transpose(0, 1).unwrap()).map_err(\|e\| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))
	}

	#[pymodule]
	fn nucleo_rust(m: &Bound<'_, PyModule>) -> PyResult<()> {
	pyo3_log::init();
	m.add_function(wrap_pyfunction!(versao, m)?)?;
	m.add_function(wrap_pyfunction!(info_build, m)?)?;
	m.add_function(wrap_pyfunction!(lint_config, m)?)?;
	m.add_function(wrap_pyfunction!(treinar_fase_a, m)?)?;
	m.add_function(wrap_pyfunction!(treinar_fase_b, m)?)?;
	m.add_function(wrap_pyfunction!(gerar_amostras, m)?)?;
	m.add_function(wrap_pyfunction!(verificar_mascara, m)?)?;
	m.add_function(wrap_pyfunction!(verificar_stiefel, m)?)?;
	m.add_class::<BpePy>()?;
	m.add_class::<ConfigPy>()?;
	m.add_class::<PonteDuplaPy>()?;
	#[cfg(feature = "ponte_async")]
	m.add_class::<PonteAsyncPy>()?;
	Ok(())
	}