Blame for scripts/31_teste_v21.py

# -*- coding: utf-8 -*-
"""31_teste_v21.py — TESTE REAL v21 (KHMAMBAJEPA, doc 33 §10).

Requisitos do usuário atendidos aqui:
  • "treinar com datasets reais verificando progresso" — corpus REAL
    (cache_dados/corpus_v16.jsonl — 1015 registros coletados na v19 de
    LeRobot/GSM8K/MATH-500/hendrycks/OpenR1/IFM; coleta fresca opcional);
  • A/B PAREADO EXATO (Teo 31.7): mesma semente, re-seed APÓS construção,
    mesmos lotes — UMA variável por braço;
  • valida as 3 melhorias v21 ISOLADAS: normalização espectral do NS,
    hot-paths Cython (equivalência Δ=0 + wallclock), tabela v17 ativa.

Braços (UMA variável por processo):
  A) baseline_v20  — NS Frobenius (patch v20) + hot-paths OFF + v17 OFF
  B) v21_ns        — NS espectral (v21) + hot-paths OFF
  C) v21_full      — NS espectral + hot-paths Cython ON (equivalência
                     Δ=0 esperada no TREINO; wallclock comparado)
  D) v17_ativo     — v21_full + treino.autorreg_v17.ativo=true

Uso (RAM 4 GiB / processos morrem entre chamadas — retomada por partes):
  python3 scripts/31_teste_v21.py --variante baseline_v20 --parte 12
  python3 scripts/31_teste_v21.py --variante v21_ns      --parte 12
  python3 scripts/31_teste_v21.py --variante v21_full    --parte 12
  python3 scripts/31_teste_v21.py --variante v17_ativo   --parte 12
  python3 scripts/31_teste_v21.py --relatorio
"""
from __future__ import annotations

import argparse
import base64
import json
import os
import random
import sys
import time

RAIZ = os.environ.get("KHTST_RAIZ", os.path.dirname(
    os.path.dirname(os.path.abspath(__file__))))
sys.path.insert(0, os.path.join(RAIZ, "src"))
sys.path.insert(0, os.path.join(RAIZ, "src", "khtst", "acelerado"))

CORPUS = os.path.join(RAIZ, "cache_dados", "corpus_v16.jsonl")
TOKENIZADOR = os.path.join(RAIZ, "cache_dados", "tokenizador_v16.json")
TOK_ANTIGO = os.path.join(RAIZ, "cache_dados", "tokenizador.json")
TELEMETRIA = os.path.join(RAIZ, "telemetria_out")
SEED = 2026
PASSOS_TOTAIS = 48


def carregar_registros(caminho: str) -> list[dict]:
    registros = []
    if not os.path.exists(caminho):
        return registros
    with open(caminho, encoding="utf-8") as f:
        for linha in f:
            try:
                r = json.loads(linha)
            except Exception:
                continue
            if isinstance(r.get("imagem"), str):
                try:
                    r["imagem"] = base64.b64decode(r["imagem"])
                except Exception:
                    r["imagem"] = None
            if isinstance(r.get("audio"), str):
                try:
                    r["audio"] = base64.b64decode(r["audio"])
                except Exception:
                    r["audio"] = None
            registros.append(r)
    return registros


def est_padrao(seed: int) -> dict:
    import numpy as np
    import torch
    torch.manual_seed(seed)
    random.seed(seed)
    np.random.seed(seed)
    return {"torch": torch.get_rng_state(), "random": random.getstate(),
            "numpy": np.random.get_state()}


def _restaurar_rng(est: dict) -> None:
    import numpy as np
    import torch
    torch.set_rng_state(est["torch"])
    random.setstate(est["random"])
    np.random.set_state(est["numpy"])


# v20 exato — braço A (replicação VERBATIM da normalização Frobenius;
# a única diferença da v21 é o bloco de normalização — doc 33 §3)
def _ns_frob_v20(G, iters=5, limiar_f32=1e-3, teto_espectral=None):
    import torch as _t
    a, b, c = 3.4445, -4.7750, 2.0315
    X = G.float() if float(G.float().norm()) < limiar_f32 else G.bfloat16()
    if G.size(-2) > G.size(-1):
        X = X.mT
    X = X / (X.norm(dim=(-2, -1), keepdim=True) + 1e-7)
    for _ in range(iters):
        A = X @ X.mT
        B = b * A + c * (A @ A)
        X = a * X + B @ X
    if G.size(-2) > G.size(-1):
        X = X.mT
    return X


def rodar_variante(nome: str, parte: int) -> dict:
    import gc
    import numpy as np
    import torch
    if not (torch.cuda.is_available() or torch.xpu.is_available()):
        torch.set_num_threads(int(os.environ.get("OMP_NUM_THREADS", "2")))
    import khtst.treino.otimizadores as O
    from khtst.config import Config
    from khtst.dados.tokenizador import TokenizadorKHTST
    from khtst.memoria.orquestrador import OrquestradorSOM
    from khtst.nucleo.modelo import KHTSTModel
    from khtst.telemetria.hub import TelemetryHub
    from khtst.treino.treinador import TreinadorExtensao
    from khtst.mamba3 import kernel as k3
    from khtst.percepcao import moe as moe_mod

    cfg = Config()
    cfg.dados.semente = SEED
    cfg.treino.warmup = 8
    cfg.treino.lote = 4
    cfg.treino.lote_efetivo = 1
    cfg.treino.qat_ultimos_passos = 0
    cfg.treino.ensemble["a_cada"] = 24
    cfg.modelo.escalacao["ativo"] = False

    # ---- A VARIÁVEL DO BRAÇO ------------------------------------------------
    if nome == "baseline_v20":
        O.newton_schulz5 = _ns_frob_v20          # NS Frobenius (v20 exato)
        k3.set_conv_cython(False)
        moe_mod.set_topk_cython(False)
        cfg.treino.autorreg_v17 = {"ativo": False}
    elif nome == "v21_ns":
        pass                                     # NS v21 default; resto OFF
    elif nome == "v21_full":
        pass                                     # defaults v21 (NS + Cython)
    elif nome == "v17_ativo":
        cfg.treino.autorreg_v17 = {"ativo": True, "a_cada": 5}
    else:
        raise ValueError(nome)

    os.makedirs(TELEMETRIA, exist_ok=True)
    hub = TelemetryHub(os.path.join(TELEMETRIA, f"treino_v21_{nome}.jsonl"))
    tk = TokenizadorKHTST(TOKENIZADOR if os.path.exists(TOKENIZADOR)
                          else TOK_ANTIGO)
    registros = carregar_registros(CORPUS)
    modelo = KHTSTModel(cfg, usar_multimodal=True)
    orquestrador = OrquestradorSOM(cfg.modelo.d_modelo, cfg.som, hub=hub,
                                   cfg_atencao=None)
    treinar = TreinadorExtensao(cfg, modelo, tk, hub, registros,
                                orquestrador_som=orquestrador)
    # PAREAMENTO EXATO (Teo 31.7): re-seed APÓS a construção
    _restaurar_rng(est_padrao(SEED))

    ckpt = os.path.join(TELEMETRIA, f"ckpt_v21_{nome}.pt")
    perdas: list[float] = []
    passo_inicial = 0
    t_acum = 0.0
    if parte and os.path.exists(ckpt):
        est = torch.load(ckpt, map_location="cpu", weights_only=False)
        modelo.load_state_dict(est["modelo"])
        treinar.otimizador.load_state_dict(est["otimizador"])
        treinar.passo_global = est["passo_global"]
        perdas = est["perdas"]
        t_acum = est.get("t_acum", 0.0)
        torch.set_rng_state(est["rng_torch"])
        random.setstate(est["rng_py"])
        np.random.set_state(est["rng_np"])
        passo_inicial = est["passo_global"]
        print(f"[retomada] passo={passo_inicial} perdas={len(perdas)}",
              flush=True)

    alvo = min(PASSOS_TOTAIS, passo_inicial + parte)
    t0 = time.time()
    while treinar.passo_global < alvo:
        perda = treinar._passo()
        if perda == perda:
            perdas.append(perda)
        if treinar.passo_global % 12 == 0:
            print(f"[{nome} passo {treinar.passo_global} perda {perda:.3f}]",
                  flush=True)
    t_acum += time.time() - t0

    resultado = {"nome": nome, "passos": treinar.passo_global,
                 "perdas": perdas, "t_acum_s": round(t_acum, 2),
                 "autorreg_v17": treinar.autorreg_v17 is not None}
    if parte and treinar.passo_global < PASSOS_TOTAIS:
        torch.save({"modelo": modelo.state_dict(),
                    "otimizador": treinar.otimizador.state_dict(),
                    "passo_global": treinar.passo_global,
                    "perdas": perdas, "t_acum": t_acum,
                    "rng_torch": torch.get_rng_state(),
                    "rng_py": random.getstate(),
                    "rng_np": np.random.get_state()}, ckpt)
    elif os.path.exists(ckpt):
        os.remove(ckpt)                        # braço completo — ckpt fora
    gc.collect()
    return resultado


def relatorio() -> None:
    res = {}
    for nome in ("baseline_v20", "v21_ns", "v21_full", "v17_ativo"):
        cam = os.path.join(TELEMETRIA, f"resultado_v21_{nome}.json")
        if os.path.exists(cam):
            res[nome] = json.load(open(cam, encoding="utf-8"))
    if len(res) < 2:
        print("resultados insuficientes — rode os 4 braços primeiro")
        return
    print("=" * 74)
    print(f"{'braço':14s} {'perda inicial':>13s} {'perda final':>11s} "
          f"{'ganho ×':>8s} {'ΔCE pareado':>12s} {'s/ passo':>9s}")
    base = res.get("baseline_v20")
    linhas = {}
    for nome, r in res.items():
        p = r["perdas"]
        ini = sum(p[:4]) / 4
        fim = sum(p[-4:]) / 4
        ganho = ini / max(fim, 1e-9)
        dce = (sum(p[-4:]) - sum(base["perdas"][-4:])) / 4 \
            if base and nome != "baseline_v20" else 0.0
        spp = r["t_acum_s"] / max(1, len(p))
        print(f"{nome:14s} {ini:13.3f} {fim:11.3f} {ganho:8.3f} "
              f"{dce:+12.4f} {spp:9.2f}")
        linhas[nome] = {"ini": round(ini, 4), "fim": round(fim, 4),
                        "ganho": round(ganho, 4), "dce": round(dce, 4),
                        "s_por_passo": round(spp, 2)}
    print("=" * 74)
    # progresso por quarto (verificação de progresso contínuo)
    for nome, r in res.items():
        p = r["perdas"]
        q = max(1, len(p) // 4)
        quarts = [round(sum(p[i * q:(i + 1) * q]) / q, 3)
                  for i in range(4) if p[i * q:(i + 1) * q]]
        print(f"{nome}: {' → '.join(map(str, quarts))}")
    out = os.path.join(RAIZ, "resultados_v21.json")
    json.dump({"braços": linhas}, open(out, "w", encoding="utf-8"),
              ensure_ascii=False, indent=1)
    print(f"salvo em {out}")


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--variante", choices=["baseline_v20", "v21_ns",
                                           "v21_full", "v17_ativo"])
    ap.add_argument("--parte", type=int, default=0,
                    help="passos NESTA chamada (retomada por ckpt)")
    ap.add_argument("--relatorio", action="store_true")
    args = ap.parse_args()
    if args.relatorio:
        relatorio()
        return
    r = rodar_variante(args.variante, args.parte)
    cam = os.path.join(TELEMETRIA, f"resultado_v21_{args.variante}.json")
    if args.parte == 0 or r["passos"] >= PASSOS_TOTAIS:
        json.dump(r, open(cam, "w", encoding="utf-8"), ensure_ascii=False)
    print(f"[{args.variante}] passos={r['passos']} "
          f"perda={r['perdas'][-1] if r['perdas'] else None}")


if __name__ == "__main__":
    main()