Blame for src/python/treino_fases.rs

//! `python::treino_fases` — extraído do monólito na v0.10.0
//! (divisão SRP, crítica Qwen3.8 lote 3 nº 1; r-35 §2).

use super::*;

use super::nucleo::ConfigPy;


/// 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))]
pub(crate) 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 = callback_treino_de_py(callback);
    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))]
pub(crate) 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 = callback_treino_de_py(callback);
    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())
    })
}


pub(crate) 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))]
pub(crate) 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}"))
}


pub(crate) 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(())
}


/// **FASE C — treino conjunto Pensamento↔backbone** (dados reais).
#[pyfunction]
pub(crate) fn treinar_fase_c(
    cfg: &ConfigPy,
    sequencias: Vec<Vec<u32>>,
    caminho_tokenizador: &str,
    passos: usize,
    caminho_gerador: Option<String>,
    caminho_telemetria: Option<String>,
    callback: Option<Bound<'_, PyAny>>,
    teto_seg: Option<f64>,
) -> PyResult<PyObject> {
    let cb = callback_treino_de_py(callback);
    let mut hub = telemetria::TelemetryHub::novo(caminho_telemetria);
    let dados = treino::DadosTreino::novo(sequencias, cfg.cfg.modelo.comprimento_ctx);
    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_c(&cfg.cfg, &tok, &dados, &mut hub, &cb, teto_seg, passos, caminho_gerador.as_deref())
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("fase C: {e}")))?;
    Python::with_gil(|py| {
        let dados = serde_json::json!({
            "resumo": {
                "passos": resumo.passos,
                "ce_inicial": resumo.ce_inicial,
                "ce_final": resumo.ce_final,
                "ce_gatilhada_sem_correcao": resumo.ce_gatilhada_sem_correcao,
                "ce_gatilhada_com_correcao": resumo.ce_gatilhada_com_correcao,
                "confianca_media_inicial": resumo.confianca_media_inicial,
                "confianca_media_final": resumo.confianca_media_final,
                "confianca_pos_correcao": resumo.confianca_pos_correcao,
                "taxa_gatilho": resumo.taxa_gatilho,
                "ece_antes": resumo.ece_antes,
                "ece_depois": resumo.ece_depois,
                "temperatura_calibrada": resumo.temperatura_calibrada,
                "tau_c": resumo.tau_c,
                "taxa_conflito_pcgrad": resumo.taxa_conflito_pcgrad,
                "decaimento_delta": resumo.decaimento_delta,
                "acordos_medios": resumo.acordos_medios,
                "escore_aprendizado": resumo.escore_aprendizado,
                "rodadas_acoplador": resumo.rodadas_acoplador,
                "timeouts_acoplador": resumo.timeouts_acoplador,
                "maior_sobreposicao_us": resumo.maior_sobreposicao_us,
            },
            "painel_v060": hub.painel_v060(),
            "telemetria": hub.dump(),
        });
        Ok(pythonize::pythonize(py, &dados)?.into())
    })
}