# -*- coding: utf-8 -*- """21_teste_v17.py — TESTE REAL v17 (doc 29): tokenizador 16384, perda TCE corrigida (JEPA normalizada + λ adaptativo), ReplaySSM, auditoria SonicMoE e torchao com EMPACOTAMENTO MoE (indexador/floresta). Variantes (UMA por processo — cada modelo completo): A) baseline_v17 — CE pura (lambda_jepa=0, replay 0) — referência de aprendizado (espera-se ~1.8× como na v16) B) v17_khmambajepa — JEPA NORMALIZADA (λ=7.5 na escala por-coordenada, Teo 29.6) + λ adaptativo pela razão PONDERADA (Teo 29.8) + ReplaySSM (S=16, Teo 29.13) + auditoria de memória I1–I4 (Teo 29.12) + quantização final com EMPACOTAMENTO MoE (Teo 29.2–29.4) + tokenizador vocab 16384 (requisito "manter vocab 16384"). Uso: python3 scripts/21_teste_v17.py --tokenizador # retreina @16384 python3 scripts/21_teste_v17.py --variante baseline_v17 --passos 120 python3 scripts/21_teste_v17.py --variante v17_khmambajepa --passos 120 [--parte N] python3 scripts/21_teste_v17.py --relatorio # consolida A/B """ from __future__ import annotations import argparse import base64 import json import os import sys RAIZ = os.environ.get("KHTST_RAIZ", "/home/z/my-project/khtst") sys.path.insert(0, os.path.join(RAIZ, "src")) CORPUS_MERGIDO = os.path.join(RAIZ, "cache_dados", "corpus_v16.jsonl") TOKENIZADOR_16K = os.path.join(RAIZ, "cache_dados", "tokenizador_v17.json") TOK_ANTIGO = os.path.join(RAIZ, "cache_dados", "tokenizador_v16.json") METRICAS = os.path.join(RAIZ, "telemetria_out", "metricas_v17") NOVOS_IDS = {"allenai/fetchman-data", "endoard/grab_ball_3cam_skin", "SoSolaris/Astra_test_20261005052625", "mulligan/sim-square-narrow-c00-teleop-mixed", "pltops/evosynth-gsm8k-smoke-cerebras-gsm8k", "heatherwhite/math-collection", "HuggingFaceH4/MATH-500", "IFM/Math-Reasoning", "EleutherAI/hendrycks_math#algebra", "EleutherAI/hendrycks_math#number_theory", "open-r1/OpenR1-Math-220k", "IFM/Math-Reasoning#socratic", "3DTopia/4DNeX-10M"} 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 retreinar_tokenizador_16k() -> dict: """REQUISITO: manter vocab 16384 — o tokenizador v16 foi treinado com 8192 (desalinhado do modelo). O v17 alinha tokenizer ≡ modelo (16384): menos UNK/bytes fallback ⟹ sequências mais curtas e CE por token mais informativa (doc 29 §7).""" from khtst.dados.tokenizador import TokenizadorKHTST registros = carregar_registros(CORPUS_MERGIDO) tk = TokenizadorKHTST() def textos(): for r in registros: if r.get("texto"): yield r["texto"] if r.get("entrada"): yield r["entrada"] if r.get("saida"): yield r["saida"] info = tk.treinar(textos(), vocab=16384, salvar_em=TOKENIZADOR_16K) info["n_registros"] = len(registros) return info def rodar_variante(nome: str, v17: bool, passos: int, parte: int = 0) -> dict: import gc import torch 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 cfg = Config() cfg.dados.semente = 2026 cfg.treino.warmup = 8 cfg.treino.lote = 4 cfg.treino.qat_ultimos_passos = 0 cfg.treino.ensemble["a_cada"] = 16 cfg.modelo.escalacao["ativo"] = False if v17: # v17 — correção da perda TCE alta (doc 29 §4): cfg.treino.autorreg["ativo"] = True cfg.treino.autorreg["a_cada"] = 5 cfg.treino.autorreg["passos_espera"] = 10 cfg.treino.autorreg["a0"] = 0.05 cfg.treino.autorreg["tau_a"] = 400.0 cfg.treino.jepa["lambda_jepa"] = 7.5 # escala por-coordenada cfg.treino.jepa["warmup_passo"] = 30 cfg.treino.jepa["vicreg_lambda"] = 0.1 cfg.treino.jepa["vicreg_gamma"] = 1.0 cfg.treino.jepa["normalizar_dim"] = True # Teo 29.6 cfg.treino.jepa["taxa_theta"] = 0.5 # Teo 29.7 # ReplaySSM (doc 29 §6): replay de estados por segmento cfg.modelo.mamba3["replay_segmento"] = 16 else: cfg.treino.autorreg["ativo"] = False cfg.treino.jepa["lambda_jepa"] = 0.0 cfg.modelo.mamba3["replay_segmento"] = 0 os.makedirs(os.path.join(RAIZ, "telemetria_out"), exist_ok=True) hub = TelemetryHub(os.path.join(RAIZ, "telemetria_out", f"treino_v17_{nome}.jsonl")) tok_path = TOKENIZADOR_16K if os.path.exists(TOKENIZADOR_16K) else TOK_ANTIGO tk = TokenizadorKHTST(tok_path) registros = carregar_registros(CORPUS_MERGIDO) 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) ckpt = os.path.join(RAIZ, "telemetria_out", f"ckpt_v17_{nome}.pt") perdas: list[float] = [] passo_alvo = (parte if parte else passos) 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"] print(f"[retomada] passo_global={est['passo_global']} " f"perdas={len(perdas)}") while treinar.passo_global < passo_alvo and treinar.passo_global < passos: perda = treinar._passo() if perda == perda: perdas.append(perda) if parte and (treinar.passo_global % 10 == 0): torch.save({"modelo": modelo.state_dict(), "otimizador": treinar.otimizador.state_dict(), "passo_global": treinar.passo_global, "perdas": perdas}, ckpt) primeiros = perdas[:8] ultimos = perdas[-8:] res = { "variante": nome, "v17": v17, "passos": len(perdas), "lote": cfg.treino.lote, "vocab_tokenizador": tk.vocab_size, "perda_media_inicial": round(sum(primeiros) / len(primeiros), 4), "perda_media_final": round(sum(ultimos) / len(ultimos), 4), "melhoria_x": round((sum(primeiros) / len(primeiros)) / max(sum(ultimos) / len(ultimos), 1e-9), 3), "perda_min": round(min(perdas), 4), "perdas_finitas": all(p == p for p in perdas), "n_registros_corpus": len(registros), } if parte: torch.save({"modelo": modelo.state_dict(), "otimizador": treinar.otimizador.state_dict(), "passo_global": treinar.passo_global, "perdas": perdas}, ckpt) if v17: res["autorreg"] = treinar.autorreg.telemetria() tel_j = getattr(treinar, "ultima_tel_jepa", None) if tel_j: res["jepa_vicreg"] = {k: v for k, v in tel_j.items() if k.startswith("jepa/")} aud = getattr(treinar, "ultima_auditoria_sonic", None) if aud: res["auditoria_sonic"] = {k: v for k, v in aud.items() if k.startswith("I") or k == "ok"} # ---- TORCHAO pós-treino COM EMPACOTAMENTO MoE (doc 29 §§2–3) ---- from khtst.quanta.torchao_quant import quantizar_torchao perda_fp32 = _perda_val(treinar, cfg, tk, registros) res["perda_val_fp32"] = perda_fp32 ev = quantizar_torchao(modelo, dict(cfg.modelo.torchao) | { "ativo": True, "empacotar_moe": True, "n_blocos_floresta": 2}) res["torchao"] = {k: (round(v, 6) if isinstance(v, float) else v) for k, v in ev.items() if k not in ("fallbacks", "moe_empacotado")} res["moe_empacotado"] = { "ok": ev["moe_empacotado"].get("ok"), "n_moes": ev["moe_empacotado"].get("n_moes"), "respeita_cota": all( v.get("respeita_cota", False) for m in ev["moe_empacotado"].get("moes", {}).values() for v in m.get("matrizes", {}).values())} perda_int8 = _perda_val(treinar, cfg, tk, registros) res["perda_val_int8"] = perda_int8 res["torchao_degradacao_perda"] = round(perda_int8 - perda_fp32, 5) else: res["perda_val_fp32"] = _perda_val(treinar, cfg, tk, registros) del treinar, modelo, orquestrador, hub gc.collect() return res def _perda_val(treinador, cfg, tk, registros, n_lotes: int = 8) -> float: """Perda em registros NOVOS (held-out das fontes v16/v17).""" import random import torch modelo = treinador.modelo modelo.eval() novos = [r for r in registros if r.get("fonte") in NOVOS_IDS or (r.get("meta") or {}).get("dataset") in NOVOS_IDS] if len(novos) < 4: novos = registros[-64:] rng = random.Random(2026) rng.shuffle(novos) perdas = [] L = cfg.modelo.comprimento_ctx with torch.no_grad(): for i in range(n_lotes): pedaco = novos[i * 2:(i + 1) * 2] if len(pedaco) < 2: break seqs = [] for s in pedaco: inp = tk.encode(s.get("entrada") or s.get("texto") or "x", tarefa=s.get("tarefa"), max_len=L // 2) rest = max(L - len(inp) - 1, 8) out = tk.encode(s.get("saida") or s.get("texto") or "y", tarefa=None, max_len=rest, com_bos=False) seqs.append((inp + out)[:L]) Tmax = max(len(s) for s in seqs) ids = torch.zeros(len(seqs), Tmax, dtype=torch.long) for k, s in enumerate(seqs): ids[k, :len(s)] = torch.tensor(s) ids[k, len(s):] = 0 alvo = ids.clone() logits, perda = modelo(ids, alvo=alvo, tarefa=pedaco[0]["tarefa"]) perdas.append(float(perda)) modelo.train() return round(sum(perdas) / max(len(perdas), 1), 4) def relatorio() -> int: a = b = None for nome in ("baseline_v17", "v17_khmambajepa"): cam = METRICAS + f"_{nome}.json" if os.path.exists(cam): with open(cam) as f: if nome == "baseline_v17": a = json.load(f) else: b = json.load(f) if not a or not b: print("Faltam métricas — rode as duas variantes primeiro.") return 1 print("=== TESTE REAL v17 — A/B PAREADO (corpus 903, vocab 16384) ===") for k in ("perda_media_inicial", "perda_media_final", "melhoria_x", "perda_min", "perda_val_fp32"): print(f" {k:24s}: baseline {a.get(k)} | v17 {b.get(k)}") if "torchao" in b: t = b["torchao"] print(f" torchao+MoE-pack: ok={t['ok']} " f"camadas={t['n_camadas_quantizadas']}" f" compressão={t['compressao']:.2f}×" f" Δperda={b['torchao_degradacao_perla'] if 'torchao_degradacao_perla' in b else b['torchao_degradacao_perda']}") print(f" moe_empacotado: {b.get('moe_empacotado')}") if "jepa_vicreg" in b: j = b["jepa_vicreg"] print(f" jepa v17: perda={j.get('jepa/perda')} " f"(v16 era 24.96) taxa={j.get('jepa/taxa_sucesso')} " f"erro_rel={j.get('jepa/erro_rel_medio')}") if "auditoria_sonic" in b: print(f" auditoria I1–I4: {b['auditoria_sonic']}") ok = (a["perdas_finitas"] and b["perdas_finitas"] and b["perda_media_final"] < b["perda_media_inicial"] and a["perda_media_final"] < a["perda_media_inicial"]) # ganho v17: o gap de melhoria entre v17 e baseline deve FECHAR vs v16 # (v16: 1.296× vs 1.811× — razão 0.716; v17 alvo: razão ≥ 0.80) try: razao_ganho = b["melhoria_x"] / a["melhoria_x"] print(f" razão de ganho v17/baseline: {razao_ganho:.3f} " f"(v16 era 0.716 — quanto mais perto de 1, menor o preço)") except Exception: pass print("TREINO CURTO v17:", "APROVADO ✓" if ok else "REPROVADO ✗") return 0 if ok else 1 def main() -> int: ap = argparse.ArgumentParser() ap.add_argument("--tokenizador", action="store_true") ap.add_argument("--variante", choices=("baseline_v17", "v17_khmambajepa")) ap.add_argument("--passos", type=int, default=120) ap.add_argument("--parte", type=int, default=0) ap.add_argument("--relatorio", action="store_true") args = ap.parse_args() if args.tokenizador: print(json.dumps(retreinar_tokenizador_16k(), ensure_ascii=False)) return 0 if args.variante: res = rodar_variante(args.variante, args.variante == "v17_khmambajepa", args.passos, parte=args.parte) os.makedirs(os.path.dirname(METRICAS), exist_ok=True) with open(METRICAS + f"_{args.variante}.json", "w") as f: json.dump(res, f, ensure_ascii=False, indent=2) print(json.dumps(res, ensure_ascii=False, indent=2)[:2500]) return 0 if args.relatorio: return relatorio() ap.print_help() return 0 if __name__ == "__main__": sys.exit(main())