scripts/tests/run_fase2_from_partial.py

#!/usr/bin/env python3
"""run_fase2_from_partial.py — Carrega estado parcial da FASE1 e executa FASE2.

V6.5-V4-canonical-256: Usa o estado parcial salvo pela FASE1 (7000+ samples)
e executa FASE2 (TREINAMENTO COM PUNIÇÃO) sobre BrunoN-Dev/corpus-ptbr-v1.

User requirement:
  "FASE2 TREINAMENTO (meta mínima 2000 samples ou mais) COM PUNIÇÃO ATIVA
   para 'BrunoN-Dev/corpus-ptbr-v1' de 100 em 100 samples"
  "LEMBRANDO que agora FASE1 e FASE2 estão treinadas no mesmo estado do modelo"
"""
import os
import sys
import time
import gc
import json
import logging
import traceback
from pathlib import Path
from datetime import datetime

# Paths
PROJECT_ROOT = Path("/home/z/my-project")
BIGRU_ROOT = PROJECT_ROOT / "BiGRU_T_version"
SRC_ROOT = BIGRU_ROOT / "src"
sys.path.insert(0, str(SRC_ROOT))

# Ambiente anti-OOM
os.environ["HF_DATASETS_DISABLE_IN_MEMORY_CACHE"] = "1"
os.environ["DATASETS_FINGERPRINT_CACHING_DISABLED"] = "1"
os.environ["TOKENIZERS_PARALLELISM"] = "false"
os.environ["HF_HUB_DISABLE_TELEMETRY"] = "1"
os.environ["HF_DATASETS_CACHE"] = "/tmp/hf_datasets_cache_v65"
os.environ["V65_ENABLE_STREAMING"] = "1"
os.environ["OMP_NUM_THREADS"] = "2"
os.environ["MKL_NUM_THREADS"] = "2"
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128,expandable_segments:True"

logging.basicConfig(
    level=logging.INFO,
    format="[%(asctime)s] [%(levelname)s] %(message)s",
    datefmt="%H:%M:%S",
    handlers=[
        logging.StreamHandler(sys.stdout),
        logging.FileHandler(str(PROJECT_ROOT / "logs" / "fase2_v4.log")),
    ],
)
logger = logging.getLogger(__name__)

from bigru_t.utils.xeon_runtime import optimize_xeon_environment
optimize_xeon_environment(verbose=False)

import torch
from bigru_t.model.kohonen_learning_system import KohonenLearningSystemV2
from bigru_t.utils.oom_guard import OomGuard
from bigru_t.data.streaming_datasets import stream_dataset
from bigru_t.model.som_metrics import compute_all_metrics

OOM_GUARD = OomGuard(max_rss_mb=2200, warn_rss_mb=1800, check_interval=2.0)
OOM_GUARD.start()

# Canônicos V6.5-V4
SOM_GRID = (4, 4, 4, 4)
HIDDEN_DIM = 1024
VOCAB_SIZE = 16384
N_HYPOTHESES = 16
MAX_N_HYPOTHESES = 32
HYP_TRAIN_STEPS = 30
HYP_HIDDEN_DIM = 256
PUNICAO_DATASET = "BrunoN-Dev/corpus-ptbr-v1"
META_MINIMA_PUNICAO = 2000
STREAM_BATCH_SIZE = 100
BATCH_SIZE = 16

HF_TOKEN = os.environ.get("HF_TOKEN")


def mem_mb() -> float:
    try:
        with open("/proc/self/status") as f:
            for line in f:
                if line.startswith("VmRSS:"):
                    return int(line.split()[1]) / 1024
    except Exception:
        pass
    return 0.0


def main() -> int:
    logger.info("=" * 80)
    logger.info("[FASE2] V6.5-V4-canonical-256 — TREINAMENTO COM PUNIÇÃO")
    logger.info("=" * 80)
    logger.info(f"  SOM grid: {SOM_GRID} (256 neurons) | HIDDEN={HIDDEN_DIM} | VOCAB={VOCAB_SIZE}")
    logger.info(f"  n_hyp: {N_HYPOTHESES}/{MAX_N_HYPOTHESES} | hyp_steps={HYP_TRAIN_STEPS}")
    logger.info(f"  Dataset: {PUNICAO_DATASET} | Meta: ≥{META_MINIMA_PUNICAO}")
    logger.info(f"  MEM start: {mem_mb():.0f}MB")
    logger.info("=" * 80)

    # 1. KLS
    logger.info("[FASE2] Inicializando KLS V2...")
    kls = KohonenLearningSystemV2(
        vocab_size=VOCAB_SIZE, hidden_dim=HIDDEN_DIM, seq_len=8,
        som_grid=SOM_GRID, alpha0=0.5, sigma0=2.0,
        lambda_ewc=0.02, N_start=10, dim_choice="y",
        hypothesis_hidden=[512, 256, 128, 64, 32, 16, 8],
        T_max=10000,
        enable_vqvae2=True, enable_reasoning=False, enable_w8a8=False,
        vqvae2_code_dim=16, vqvae2_num_codes_top=64, vqvae2_num_codes_bot=128,
        enable_attention=True, attention_n_heads=8,
        n_hypotheses=N_HYPOTHESES, max_n_hypotheses=MAX_N_HYPOTHESES,
        min_n_hypotheses=4, n_trials=3, min_n_trials=1, max_n_trials=6,
        hyp_train_steps=HYP_TRAIN_STEPS, min_hyp_train_steps=10, max_hyp_train_steps=80,
        hyp_lr=1e-4, hyp_hidden_dim=HYP_HIDDEN_DIM,
        loss_history_window=8, punishment_window=12,
    )
    logger.info(f"[FASE2] KLS V2 init: {mem_mb():.0f}MB")
    kls.buffer_max_size = 128  # OOM-safety

    # 2. Tokenizer
    corpus_inicial = [
        "o gato dorme na cama", "a casa eh azul", "ele corre rapido",
        "ela canta uma musica", "o sol nasceu hoje", "nos vamos viajar",
        "o livro esta na mesa", "a menina brinca no parque",
        "ola como voce esta", "qual e o seu nome",
        "calcule dois mais dois", "traduza hello para portugues",
        "instrucao para resolver o problema", "resposta para a pergunta",
        "luva de pedreiro tavila", "lula reserva valor",
        "amazonas forca tarefa vitimas",
    ]
    kls.tokenizer.fit(corpus_inicial)
    logger.info(f"[FASE2] Tokenizer fitted: {mem_mb():.0f}MB")

    # 3. Carrega estado parcial (procura o mais recente)
    partials = sorted(BIGRU_ROOT.glob("v6_5_v2_conhecimento_partial_d*.pt"))
    if not partials:
        logger.error("[FASE2] Nenhum estado parcial encontrado. Abortando.")
        return 1
    partial_state_path = partials[-1]
    logger.info(f"[FASE2] Carregando estado: {partial_state_path.name}")

    try:
        state = torch.load(str(partial_state_path), map_location="cpu", weights_only=False)
        meta = state.get("_meta", {})
        logger.info(f"[FASE2] Estado: phase={meta.get('phase')}, "
                    f"samples={meta.get('total_samples')}, step={meta.get('step')}")

        if "som_weights" in state:
            kls.som.weights.data.copy_(state["som_weights"])
        if "embedding_state" in state:
            kls.embedding.load_state_dict(state["embedding_state"])
        if "hypothesis_ensemble_state" in state:
            try:
                kls.hypothesis_ensemble.load_state_dict(state["hypothesis_ensemble_state"])
            except Exception as e:
                logger.warning(f"[FASE2] hyp_ensemble load failed: {e}")
        if "delta_scale" in state:
            try:
                kls.delta_scale.data.copy_(state["delta_scale"])
            except Exception:
                pass
        if "label_registry" in state:
            try:
                kls.label_registry = state["label_registry"]
            except Exception:
                pass

        kls.time_counter = meta.get("time_counter", 7000)
        kls.training_ready = True
        fase1_samples = meta.get("total_samples", 7000)
        logger.info(f"[FASE2] Estado carregado: {mem_mb():.0f}MB, fase1_samples={fase1_samples}")
    except Exception as e:
        logger.error(f"[FASE2] Falha ao carregar estado: {e}")
        traceback.print_exc()
        return 1

    # 4. FASE2 — streaming + process_batch_v2 com punição
    logger.info("\n[FASE2] Iniciando TREINAMENTO COM PUNIÇÃO...")
    total_samples = 0
    step = 0
    punishment_events = []
    hypotheses_trainings = []
    delta_applications = []
    som_metrics_log = []
    t_start = time.time()

    try:
        sample_iter = stream_dataset(
            dataset_name=PUNICAO_DATASET,
            max_samples=META_MINIMA_PUNICAO,
            hf_token=HF_TOKEN,
        )

        chunk_buffer = []
        chunk_idx = 0

        for sample in sample_iter:
            chunk_buffer.append(sample)
            if len(chunk_buffer) >= STREAM_BATCH_SIZE:
                chunk_idx += 1
                chunk_texts = [
                    s.raw_text if hasattr(s, "raw_text") else str(s)
                    for s in chunk_buffer
                ]
                total_samples += len(chunk_texts)
                logger.info(f"[FASE2] Chunk {chunk_idx}: {len(chunk_texts)} samples "
                            f"(total={total_samples}/{META_MINIMA_PUNICAO}), MEM={mem_mb():.0f}MB")

                # Labels binários determinísticos baseados em hash
                labels = [hash(s) % 2 for s in chunk_texts]

                # Processa em sub-batches
                for bs in range(0, len(chunk_texts), BATCH_SIZE):
                    batch_sents = chunk_texts[bs: bs + BATCH_SIZE]
                    batch_labels = labels[bs: bs + BATCH_SIZE]
                    try:
                        result = kls.process_batch_v2(
                            batch_sents, batch_labels,
                            dataset_name=PUNICAO_DATASET,
                            enable_punishment=True,
                        )
                        step += 1
                        action = result.get("action", "none")
                        if action != "none":
                            logger.info(f"[FASE2] Step {step}: action={action}, "
                                        f"acc={result.get('accuracy', 0):.3f}")
                            if "train_hypotheses" in action:
                                hyp_info = result.get("hypotheses_training", {})
                                hypotheses_trainings.append({
                                    "step": step,
                                    "loss_final": hyp_info.get("loss_final"),
                                })
                            elif "apply_best_delta" in action:
                                delta_info = result.get("delta_application", {})
                                delta_applications.append({
                                    "step": step,
                                    "best_acc": delta_info.get("best_acc"),
                                })
                            punishment_events.append({
                                "step": step,
                                "action": action,
                                "accuracy": result.get("accuracy", 0),
                            })
                    except (MemoryError, RuntimeError) as oom_err:
                        is_oom = (
                            isinstance(oom_err, MemoryError)
                            or "out of memory" in str(oom_err).lower()
                        )
                        if is_oom:
                            logger.error(f"[FASE2] OOM step {step}: {str(oom_err)[:200]}")
                            gc.collect(); gc.collect()
                            time.sleep(2)
                            continue
                        raise

                    if step % 4 == 0:
                        gc.collect()
                    time.sleep(0.3)

                # Métricas SOM após chunk
                try:
                    buf = kls.buffer_4d[-64:] if kls.buffer_4d else []
                    som_metrics = compute_all_metrics(kls.som, buf)
                    som_metrics_log.append({
                        "chunk": chunk_idx,
                        "total_samples": total_samples,
                        "qe": float(som_metrics.get("quantization_error", 0)),
                        "te": float(som_metrics.get("topological_error", 0)),
                        "kl": float(som_metrics.get("kaski_lagus_error", 0)),
                        "ve": float(som_metrics.get("explained_variance_share", 0)),
                        "dead_rate": float(som_metrics.get("dead_neuron_rate", {}).get("dead_neuron_rate", 0)),
                    })
                    logger.info(
                        f"[FASE2] SOM: QE={som_metrics_log[-1]['qe']:.4f}, "
                        f"TE={som_metrics_log[-1]['te']:.4f}, "
                        f"KL={som_metrics_log[-1]['kl']:.4f}, "
                        f"VE={som_metrics_log[-1]['ve']:.4f}, "
                        f"dead={som_metrics_log[-1]['dead_rate']:.3f}"
                    )
                except Exception as e:
                    logger.warning(f"[FASE2] Métricas SOM falharam: {e}")

                # Salva estado parcial
                try:
                    partial_path = BIGRU_ROOT / f"v6_5_v2_punicão_partial_c{chunk_idx}.pt"
                    for old in BIGRU_ROOT.glob("v6_5_v2_punicão_partial_c*.pt"):
                        if old != partial_path:
                            old.unlink(missing_ok=True)
                    torch.save({
                        "_meta": {
                            "reason": f"punicao_after_chunk_{chunk_idx}",
                            "step": step,
                            "total_samples": total_samples,
                            "timestamp": datetime.now().isoformat(),
                            "version": "V6.5-V4-canonical-256",
                            "phase": "punicao_partial",
                            "som_grid": list(SOM_GRID),
                            "n_neurons": 256,
                            "fase1_samples": fase1_samples,
                        },
                        "som_weights": kls.som.weights.data,
                        "embedding_state": kls.embedding.state_dict(),
                        "hypothesis_ensemble_state": kls.hypothesis_ensemble.state_dict(),
                        "delta_scale": kls.delta_scale.data,
                        "label_registry": kls.label_registry,
                    }, str(partial_path))
                    logger.info(f"[FASE2] Estado parcial salvo: {partial_path.name}")
                except Exception as e:
                    logger.warning(f"[FASE2] Save parcial falhou: {e}")

                gc.collect(); gc.collect()
                time.sleep(1.0)

                if total_samples >= META_MINIMA_PUNICAO:
                    logger.info(f"[FASE2] Meta atingida: {total_samples} ≥ {META_MINIMA_PUNICAO}")
                    break

                chunk_buffer = []

        # Processa chunk final se houver
        if chunk_buffer and total_samples < META_MINIMA_PUNICAO:
            chunk_idx += 1
            chunk_texts = [
                s.raw_text if hasattr(s, "raw_text") else str(s)
                for s in chunk_buffer
            ]
            total_samples += len(chunk_texts)
            logger.info(f"[FASE2] Chunk final {chunk_idx}: {len(chunk_texts)} samples "
                        f"(total={total_samples})")
            labels = [hash(s) % 2 for s in chunk_texts]
            for bs in range(0, len(chunk_texts), BATCH_SIZE):
                batch_sents = chunk_texts[bs: bs + BATCH_SIZE]
                batch_labels = labels[bs: bs + BATCH_SIZE]
                try:
                    result = kls.process_batch_v2(
                        batch_sents, batch_labels,
                        dataset_name=PUNICAO_DATASET,
                        enable_punishment=True,
                    )
                    step += 1
                    if result.get("action", "none") != "none":
                        punishment_events.append({
                            "step": step,
                            "action": result.get("action"),
                            "accuracy": result.get("accuracy", 0),
                        })
                except Exception as e:
                    logger.warning(f"[FASE2] Erro no chunk final: {e}")

    except Exception as e:
        logger.error(f"[FASE2] Erro durante FASE2: {e}")
        traceback.print_exc()

    elapsed = time.time() - t_start

    # 5. Estado final unificado
    logger.info("\n[FASE2] Salvando estado final unificado...")
    final_state_path = BIGRU_ROOT / "v6_5_v2_model_states.pt"
    try:
        torch.save({
            "_meta": {
                "reason": "end_of_training_v65_v4",
                "step": step,
                "total_samples": total_samples,
                "timestamp": datetime.now().isoformat(),
                "version": "V6.5-V4-canonical-256",
                "phase": "end_of_training",
                "som_grid": list(SOM_GRID),
                "n_neurons": 256,
                "hidden_dim": HIDDEN_DIM,
                "vocab_size": VOCAB_SIZE,
                "n_hypotheses": N_HYPOTHESES,
                "max_n_hypotheses": MAX_N_HYPOTHESES,
                "hyp_train_steps": HYP_TRAIN_STEPS,
                "hyp_hidden_dim": HYP_HIDDEN_DIM,
                "buffer_max_size": kls.buffer_max_size,
                "fase1_samples": fase1_samples,
                "fase2_samples": total_samples,
            },
            "som_weights": kls.som.weights.data,
            "embedding_state": kls.embedding.state_dict(),
            "hypothesis_ensemble_state": kls.hypothesis_ensemble.state_dict(),
            "delta_scale": kls.delta_scale.data,
            "label_registry": kls.label_registry,
        }, str(final_state_path))
        logger.info(f"[FASE2] Estado final salvo: {final_state_path}")
    except Exception as e:
        logger.error(f"[FASE2] Falha ao salvar estado final: {e}")

    # 6. Relatório
    report = {
        "version": "V6.5-V4-canonical-256",
        "timestamp": datetime.now().isoformat(),
        "config": {
            "som_grid": list(SOM_GRID), "n_neurons": 256,
            "hidden_dim": HIDDEN_DIM, "vocab_size": VOCAB_SIZE,
            "n_hypotheses": N_HYPOTHESES, "max_n_hypotheses": MAX_N_HYPOTHESES,
            "hyp_train_steps": HYP_TRAIN_STEPS, "hyp_hidden_dim": HYP_HIDDEN_DIM,
            "buffer_max_size": kls.buffer_max_size,
        },
        "fase1_summary": {
            "total_samples": fase1_samples,
            "meta_atingida": fase1_samples >= 8000,
            "state_file": partial_state_path.name,
        },
        "fase2_summary": {
            "total_samples": total_samples,
            "meta_minima": META_MINIMA_PUNICAO,
            "meta_atingida": total_samples >= META_MINIMA_PUNICAO,
            "elapsed_s": elapsed,
            "punishment_events": len(punishment_events),
            "hypotheses_trainings": len(hypotheses_trainings),
            "delta_applications": len(delta_applications),
        },
        "som_metrics_log": som_metrics_log,
        "punishment_events_last": punishment_events[-10:],
        "hypotheses_trainings_last": hypotheses_trainings[-5:],
        "delta_applications_last": delta_applications[-5:],
        "oom_guard_stats": OOM_GUARD.get_stats(),
    }
    report_path = BIGRU_ROOT / "v6_5_v4_fase2_report.json"
    with open(report_path, "w") as f:
        json.dump(report, f, indent=2, ensure_ascii=False, default=str)
    logger.info(f"[FASE2] Relatório salvo: {report_path}")

    OOM_GUARD.stop()
    del kls
    gc.collect()

    logger.info("\n" + "=" * 80)
    logger.info("[FASE2] RESUMO FINAL")
    logger.info("=" * 80)
    logger.info(f"  FASE1 samples : {fase1_samples} (meta=8000)")
    logger.info(f"  FASE2 samples : {total_samples} (meta={META_MINIMA_PUNICAO})")
    logger.info(f"  Punishments   : {len(punishment_events)}")
    logger.info(f"  Hyp trainings : {len(hypotheses_trainings)}")
    logger.info(f"  Delta applies : {len(delta_applications)}")
    logger.info(f"  Elapsed       : {elapsed:.1f}s")
    logger.info(f"  Peak RSS      : {OOM_GUARD.get_stats()['peak_rss_mb']:.0f}MB")
    logger.info(f"  State file    : {final_state_path}")
    logger.info("=" * 80)

    return 0 if total_samples > 0 else 1


if __name__ == "__main__":
    sys.exit(main())