pampar/coder/v3/modelo.py

# SPDX-License-Identifier: BUSL-1.1
# Copyright (c) 2024-2026 Lucas Ricardo Mella Chillemi
"""
PamparV3 — Modelo principal con arquitectura 2D + Mixed Selectivity.

Arquitectura:
  tok_emb [48K, 640]
    → TalamoInicial (LLAVES + attn_proj + context_conv)
        → terr_acts [B, L, 4], zona_acts [B, L, 52]
    → 4 streams inicializados desde tok_emb
    → [NivelProfundo × n_levels]
        cada nivel: attn compartida + re-routing
                  + 1 FFN compartido + 4 ContextModulator (FiLM)
                  + lateral gates (fibras blancas)
    → norm_f (RMSNorm)
    → lm_head (weight-tied con tok_emb)

Mixed Selectivity: 1 FFN compartido por nivel, modulado por 63 indicadores
contextuales (zona_acts + terr_acts + depth + conf + stream_id).
La misma memoria se lee de formas diferentes según el contexto.
Ahorro: ~32M params vs 4 FFN independientes por nivel.

4 streams × 5 niveles = grilla 2D donde:
  - La profundidad refina el significado (como capas corticales)
  - La anchura especializa por tipo de token (como áreas corticales)
  - Las lateral gates comunican áreas a cada nivel (como fibras blancas)
  - El re-routing adapta qué área lidera según el contexto acumulado
  - Los moduladores permiten mixed selectivity (misma neurona, roles múltiples)
"""

from __future__ import annotations

from typing import TYPE_CHECKING, Dict, List, Optional, Tuple

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.checkpoint

from .bloques import NivelProfundo, RMSNorm
from .config import PRESET_V3, ConfigV3
from .talamo import TalamoInicial

if TYPE_CHECKING:
    from .engrama_stream import BancoEngrama


class PamparV3(nn.Module):
    """
    PAMPAr-Coder v3: arquitectura 2D con streams especializados y profundidad.

    ~110M parámetros con PRESET_V3 (dim=640, 5 niveles, vocab=48K).
    """

    def __init__(self, config: ConfigV3 = PRESET_V3):
        super().__init__()
        self.config = config

        # Embedding de tokens (weight-tied con lm_head)
        self.tok_emb = nn.Embedding(config.vocab_size, config.dim)
        self.emb_drop = nn.Dropout(config.dropout)

        # Tálamo inicial: routing completo de tokens a zonas/territorios
        self.talamo = TalamoInicial(config)

        # Grilla 2D: n_levels niveles de profundidad
        self.niveles = nn.ModuleList(
            [NivelProfundo(config, nivel_idx=i) for i in range(config.n_levels)]
        )

        # Normalización final (antes de lm_head)
        self.norm_f = RMSNorm(config.dim)

        # LM head
        self.lm_head = nn.Linear(config.dim, config.vocab_size, bias=False)

        # Weight tying: embedding y lm_head comparten pesos
        # Ahorra vocab_size × dim parámetros (48K × 640 = 30.7M fp16)
        self.lm_head.weight = self.tok_emb.weight

        # Inicialización
        self._init_weights()

    def _init_weights(self) -> None:
        """Inicialización estilo GPT-NeoX / Llama: N(0, 0.02)."""

        def _init(module: nn.Module) -> None:
            if isinstance(module, nn.Linear):
                nn.init.normal_(module.weight, mean=0.0, std=0.02)
                if module.bias is not None:
                    nn.init.zeros_(module.bias)
            elif isinstance(module, nn.Embedding):
                nn.init.normal_(module.weight, mean=0.0, std=0.02)

        self.apply(_init)

    def registrar_tokenizer(self, tokenizer: object) -> None:
        """Registra el tokenizer en el Tálamo para que LLAVES funcione."""
        self.talamo.registrar_tokenizer(tokenizer)

    def _enable_kv_cache(self) -> None:
        """Enable KV cache mode for generation (inference only)."""
        for nivel in self.niveles:
            nivel.attn._use_kv_cache = True
            nivel.attn._kv_cache = None
            nivel.attn._start_pos = 0

    def _disable_kv_cache(self) -> None:
        """Disable KV cache and free cached tensors."""
        for nivel in self.niveles:
            nivel.attn._use_kv_cache = False
            nivel.attn._kv_cache = None
            nivel.attn._start_pos = 0

    def _set_cache_pos(self, pos: int) -> None:
        """Set the position offset for RoPE in all attention layers."""
        for nivel in self.niveles:
            nivel.attn._start_pos = pos

    def set_train_norm_clamp(self, enabled: bool) -> None:
        """Activa/desactiva norm clamping durante training en todos los niveles."""
        for nivel in self.niveles:
            nivel._train_norm_clamp = enabled

    def _combinar_streams(
        self,
        streams: List[torch.Tensor],
        terr_acts: torch.Tensor,
    ) -> torch.Tensor:
        """
        Combina los 4 streams en una representación unificada.

        Ponderada por activación territorial: el stream más activo domina.
        Normalizada para que los pesos sumen 1.

        Args:
            streams:   [n_streams × [B, L, D]]
            terr_acts: [B, L, n_streams]
        Returns:
            x: [B, L, D]
        """
        # Normalizar pesos territoriales (softmax sobre streams)
        weights = F.softmax(terr_acts, dim=-1)  # [B, L, 4]
        return sum(
            streams[t] * weights[:, :, t : t + 1] for t in range(self.config.n_streams)
        )

    def forward(
        self,
        input_ids: torch.Tensor,
        targets: Optional[torch.Tensor] = None,
        use_early_exit: bool = False,
        banco_engrama: Optional[BancoEngrama] = None,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Dict]:
        """
        Forward pass de PamparV3.

        Args:
            input_ids:      [B, L] token IDs
            targets:        [B, L] labels (-100 = ignorar)
            use_early_exit: salir antes si confianza suficiente
            banco_engrama:  banco de engramas para inyección (None = sin inyección)

        Returns:
            logits: [B, L, vocab_size]
            loss:   scalar (si targets provisto)
            info:   {'exit_nivel': int, 'terr_acts': tensor}
        """
        B, L = input_ids.shape

        # 1. Embedding
        x = self.emb_drop(self.tok_emb(input_ids))  # [B, L, D]

        # 2. Tálamo inicial: routing de tokens a zonas y territorios
        terr_acts, zona_acts = self.talamo(x, input_ids)  # [B,L,4], [B,L,52]

        # 3. Inicializar los 4 streams desde el mismo embedding
        #    Cada stream parte del mismo punto y se especializa a lo largo
        #    de los n_levels niveles
        streams: List[torch.Tensor] = [x.clone() for _ in range(self.config.n_streams)]

        # 4. Pasar por cada nivel de profundidad
        info: Dict = {"exit_nivel": self.config.n_levels, "terr_acts": terr_acts}

        for i, nivel in enumerate(self.niveles):
            if self.config.use_checkpoint and self.training and not use_early_exit:
                # Gradient checkpointing: ahorra VRAM no guardando activaciones
                # Se usan lambdas para pasar args no-tensor (agregar_fn)
                def create_checkpoint_fn(n):
                    def fn(*stream_tensors):
                        s_list = list(stream_tensors[:-2])
                        ta = stream_tensors[-2]
                        za = stream_tensors[-1]
                        new_s, new_ta, _ = n(
                            s_list,
                            ta,
                            TalamoInicial.agregar_fn,
                            zona_acts=za,
                        )
                        return (*new_s, new_ta)

                    return fn

                result = torch.utils.checkpoint.checkpoint(
                    create_checkpoint_fn(nivel),
                    *streams,
                    terr_acts,
                    zona_acts,
                    use_reentrant=False,
                )
                streams = list(result[: self.config.n_streams])
                terr_acts = result[self.config.n_streams]
                conf = 0.0  # No calculada durante checkpointing
            else:
                streams, terr_acts, conf = nivel(
                    streams,
                    terr_acts,
                    TalamoInicial.agregar_fn,
                    banco_engrama=banco_engrama,
                    zona_acts=zona_acts,
                )

            # Early Exit: si el nivel más difícil del 10% ya tiene confianza
            if use_early_exit and conf > self.config.umbral_exit:
                if i >= self.config.capas_min - 1:
                    info["exit_nivel"] = i + 1
                    break

        # 5. Combinar streams en representación final
        x_final = self._combinar_streams(streams, terr_acts)  # [B, L, D]

        # 6. Norm final + LM head
        x_final = self.norm_f(x_final)
        logits = self.lm_head(x_final)  # [B, L, vocab_size]

        # 7. Loss
        loss = None
        if targets is not None:
            loss = F.cross_entropy(
                logits.reshape(-1, self.config.vocab_size),
                targets.reshape(-1),
                ignore_index=-100,
            )

        return logits, loss, info

    @torch.no_grad()
    def generate(
        self,
        prompt_ids: torch.Tensor,
        max_tokens: int = 256,
        temperature: float = 0.8,
        top_k: int = 50,
        top_p: float = 0.95,
        banco_engrama: Optional[BancoEngrama] = None,
    ) -> torch.Tensor:
        """
        Generación autoregresiva con KV cache y nucleus sampling.

        Usa prefill (procesa todo el prompt de una sola vez) +
        decode (un token por paso, reutilizando el KV cache).

        Args:
            prompt_ids:    [1, L] prompt tokenizado
            max_tokens:    máximo tokens a generar
            temperature:   diversidad (menor = más determinista)
            top_k:         Top-K sampling (0 = desactivado)
            top_p:         Nucleus sampling (1.0 = desactivado)
            banco_engrama: banco de engramas para inyección adaptativa

        Returns:
            [1, L+N] tokens generados
        """
        self.eval()
        generated = prompt_ids.clone()
        prompt_len = prompt_ids.shape[1]

        try:
            self._enable_kv_cache()

            # --- Prefill: procesar todo el prompt, poblar KV cache ---
            self._set_cache_pos(0)
            logits, _, _ = self.forward(
                prompt_ids,
                use_early_exit=False,
                banco_engrama=banco_engrama,
            )
            logits = logits[:, -1, :] / temperature

            for _ in range(max_tokens):
                # Top-K filtering
                if top_k > 0:
                    v, _ = logits.topk(top_k)
                    logits[logits < v[:, [-1]]] = float("-inf")

                # Nucleus (Top-P) filtering
                if top_p < 1.0:
                    sorted_logits, sorted_idx = logits.sort(descending=True)
                    cumprobs = sorted_logits.softmax(dim=-1).cumsum(dim=-1)
                    remove_mask = cumprobs - sorted_logits.softmax(dim=-1) > top_p
                    sorted_logits[remove_mask] = float("-inf")
                    logits = torch.zeros_like(logits).scatter(
                        1,
                        sorted_idx,
                        sorted_logits,
                    )

                # Guard contra NaN/Inf (frecuente en early training)
                if torch.isnan(logits).any() or torch.isinf(logits).any():
                    logits = torch.nan_to_num(logits, nan=0.0, posinf=1e4, neginf=-1e4)

                probs = F.softmax(logits, dim=-1)
                if torch.isnan(probs).any() or (probs < 0).any():
                    probs = torch.ones_like(probs) / probs.shape[-1]

                next_tok = torch.multinomial(probs, 1)
                generated = torch.cat([generated, next_tok], dim=1)

                # Stop en EOS (token 0)
                if generated.shape[0] == 1 and next_tok.item() == 0:
                    break

                # --- Decode: un token a la vez, KV cache se reutiliza ---
                cur_pos = generated.shape[1] - 1
                if cur_pos >= self.config.max_seq_len:
                    break

                self._set_cache_pos(cur_pos)
                logits, _, _ = self.forward(
                    next_tok,
                    use_early_exit=False,
                    banco_engrama=banco_engrama,
                )
                logits = logits[:, -1, :] / temperature

        finally:
            self._disable_kv_cache()

        return generated

    def count_params(self) -> Dict[str, int]:
        """Cuenta parámetros por componente."""
        return {
            "embeddings": self.tok_emb.weight.numel(),
            "talamo_inicial": sum(p.numel() for p in self.talamo.parameters()),
            "niveles": sum(p.numel() for p in self.niveles.parameters()),
            "norm_f": sum(p.numel() for p in self.norm_f.parameters()),
            "total": sum(p.numel() for p in self.parameters()),
            "total_sin_embedding": sum(
                p.numel()
                for name, p in self.named_parameters()
                if "tok_emb" not in name
            ),
        }

    def describe(self) -> str:
        """Descripción legible de la arquitectura."""
        p = self.count_params()
        cfg = self.config
        mem = cfg.memory_estimate_mb()
        return (
            f"PamparV3\n"
            f"  Arquitectura:  {cfg.n_streams} streams × {cfg.n_levels} niveles (2D)\n"
            f"  Dimensiones:   dim={cfg.dim}, heads={cfg.n_heads} "
            f"(GQA {cfg.n_kv_heads} KV), ffn_hidden={cfg.ffn_hidden}\n"
            f"  Vocab:         {cfg.vocab_size:,} tokens, seq_len={cfg.max_seq_len}\n"
            f"  Parámetros:    {p['total'] / 1e6:.1f}M total "
            f"({p['total_sin_embedding'] / 1e6:.1f}M sin embedding)\n"
            f"  VRAM model:    {mem['model_fp16_mb']}MB (fp16)\n"
            f"  VRAM training: {mem['training_total_mb']}MB (model+grad+Adam)\n"
            f"  KV cache:      {mem['kv_cache_inference_mb']}MB "
            f"(batch=1, seq=4096)\n"
        )


# =============================================================================
# FACTORY
# =============================================================================


def crear_modelo_v3(config: ConfigV3 = PRESET_V3) -> PamparV3:
    """Crea un modelo PamparV3 con la configuración dada."""
    return PamparV3(config)