cnn_bigru/inference/inference.py

"""inference.py — Inferência com Temperatura, Top-K, Top-P e Penalidades.

Implementa as técnicas descritas em dados.txt (seção sobre geração avançada):
    1. Temperatura: ajusta entropia das previsões
    2. Top-K: filtra os K tokens mais prováveis
    3. Top-P (Nucleus Sampling): limite dinâmico por probabilidade acumulada
    4. Presence Penalty: pune se já apareceu
    5. Frequency Penalty: pune proporcionalmente à frequência

VERSÃO 2.0:
    - Usa GeneratorCNNBiGRU (quando disponível) para geração autoregressiva REAL
    - Integra ContextWindowManager para sequências longas (sliding window)
    - Fallback para MultimodalCNNBiGRU.lm_head quando generator não disponível
    - PPL evaluation correta usando GeneratorCNNBiGRU (não single-step broadcast)
"""
from __future__ import annotations

import logging
import math
from typing import Any, Dict, List, Optional, Tuple

import torch
import torch.nn as nn
import torch.nn.functional as F

from ..models.context_window import ContextWindowManager, ContextWindowConfig
from ..models.generator_verifier import GeneratorCNNBiGRU
from ..models.multimodal_model import MultimodalCNNBiGRU

logger = logging.getLogger(__name__)


# ============================================================================
# Aplicar penalidades, temperatura, top-k, top-p
# ============================================================================

def apply_sampling_filters(
    logits: torch.Tensor,
    token_counts: Dict[int, int],
    temperature: float,
    top_k: int,
    top_p: float,
    presence_penalty: float,
    frequency_penalty: float,
) -> Tuple[torch.Tensor, Dict[str, int]]:
    """Aplica presença/frequência penalidades + temperatura + top-k + top-p.

    Args:
        logits: [B, V] (será modificado in-place via clone)
        token_counts: dict {token_id: count} para penalidades
        temperature: ajusta entropia (0 = greedy)
        top_k: filtra K mais prováveis
        top_p: limite de probabilidade acumulada
        presence_penalty: pune token se já apareceu
        frequency_penalty: pune proporcionalmente à frequência

    Returns:
        (filtered_logits, metrics_dict)
    """
    metrics = {
        "tokens_after_topk": 0,
        "tokens_after_topp": 0,
    }
    logits = logits.clone()

    # Penalidades de repetição (direct logit modification)
    if presence_penalty != 0.0 or frequency_penalty != 0.0:
        for tok_id, count in token_counts.items():
            if 0 <= tok_id < logits.size(-1):
                logits[:, tok_id] -= (presence_penalty + count * frequency_penalty)

    # Temperatura
    if temperature > 0.0:
        logits = logits / temperature
    # Se temperature == 0, chamador deve usar greedy (argmax)

    # Top-K
    if top_k > 0 and top_k < logits.size(-1):
        k = min(top_k, logits.size(-1))
        top_vals, _ = torch.topk(logits, k, dim=-1)
        thresh = top_vals[:, -1:]
        logits[logits < thresh] = float("-inf")
        metrics["tokens_after_topk"] = k

    # Top-P (Nucleus Sampling)
    if top_p < 1.0:
        sorted_logits, sorted_idx = torch.sort(logits, descending=True, dim=-1)
        sorted_probs = F.softmax(sorted_logits, dim=-1)
        cum_probs = torch.cumsum(sorted_probs, dim=-1)

        # Remove tokens cuja probabilidade acumulada ultrapassa top_p
        # (mantém o primeiro token mesmo se já ultrapassar)
        remove = cum_probs > top_p
        remove[:, 1:] = remove[:, :-1].clone()
        remove[:, 0] = False

        indices_to_remove = remove.scatter(1, sorted_idx, remove)
        logits[indices_to_remove] = float("-inf")
        metrics["tokens_after_topp"] = int((~indices_to_remove[0]).sum().item())

    return logits, metrics


# ============================================================================
# Geração Autoregressiva (V2.0 — CORRIGIDA)
# ============================================================================

@torch.no_grad()
def generate_with_sampling(
    model: nn.Module,
    tokenizer,
    prompt_a: str,
    prompt_b: str = "",
    images: Optional[torch.Tensor] = None,
    audios: Optional[torch.Tensor] = None,
    max_new_tokens: int = 32,
    temperature: float = 0.7,
    top_k: int = 40,
    top_p: float = 0.9,
    presence_penalty: float = 0.3,
    frequency_penalty: float = 0.3,
    device: str = "cpu",
    generator: Optional[GeneratorCNNBiGRU] = None,
    context_window: Optional[ContextWindowManager] = None,
) -> Dict[str, Any]:
    """Gera texto com amostragem avançada.

    VERSÃO 2.0:
      - Se `generator` (GeneratorCNNBiGRU) for fornecido, usa geração
        autoregressiva REAL com o decoder GRU + Bahdanau attention.
      - Caso contrário, fallback para MultimodalCNNBiGRU em modo "generate"
        (single-step — NÃO é autoregressivo real, apenas para testes).
      - Se `context_window` for fornecido, aplica sliding window para
        suportar sequências longas (> max_seq_len do modelo).

    Args:
        model: MultimodalCNNBiGRU
        tokenizer: BBPETokenizer
        prompt_a, prompt_b: textos de entrada
        images, audios: tensores opcionais
        max_new_tokens: máximo de tokens a gerar
        temperature: ajusta entropia (0 = greedy)
        top_k: filtra K mais prováveis
        top_p: limite de probabilidade acumulada
        presence_penalty: pune token se já apareceu
        frequency_penalty: pune proporcionalmente à frequência
        device: dispositivo de computação
        generator: GeneratorCNNBiGRU opcional (recomendado para geração real)
        context_window: gerenciador de janela de contexto opcional

    Returns:
        dict com:
            text: texto gerado
            token_ids: lista de IDs gerados
            metrics: dicionário de métricas
            used_generator: bool indicando se usou generator real
    """
    model.eval()
    if generator is not None:
        generator.eval()

    # Tokeniza prompts
    ids_a = tokenizer.encode(prompt_a, add_special=True)
    ids_b = tokenizer.encode(prompt_b, add_special=True) if prompt_b else [tokenizer.pad_id] * len(ids_a)

    # Pad para mesmo comprimento
    max_len = max(len(ids_a), len(ids_b), 8)
    pad_id = tokenizer.pad_id
    ids_a = ids_a + [pad_id] * (max_len - len(ids_a))
    ids_b = ids_b + [pad_id] * (max_len - len(ids_b))

    input_a = torch.tensor([ids_a], dtype=torch.long, device=device)
    input_b = torch.tensor([ids_b], dtype=torch.long, device=device)

    # Prepara imagens/áudios
    if images is not None:
        images = images.unsqueeze(0).to(device) if images.dim() == 3 else images.to(device)
    if audios is not None:
        audios = audios.unsqueeze(0).to(device) if audios.dim() == 3 else audios.to(device)

    # Histórico de contagem para penalidades
    token_counts: Dict[int, int] = {}
    generated_ids: List[int] = []

    eos_id = tokenizer.eos_id
    bos_id = tokenizer.bos_id

    metrics = {
        "steps": 0,
        "fallback_to_greedy": 0,
        "tokens_after_topk": 0,
        "tokens_after_topp": 0,
        "used_generator": generator is not None,
    }

    # Inicializar context window se fornecido
    # IMPORTANTE: o context_window é usado APENAS para truncar a janela de input
    # (não para cache KV — o generator faz seu próprio forward a cada step).
    # Para evitar misturar o histórico de input_a e input_b, usamos
    # truncamento simples em vez de context_window.append_tokens (que acumula).
    max_seq = 64  # janela máxima padrão
    if context_window is not None:
        try:
            max_seq = context_window.config.max_window
        except AttributeError:
            pass

    # Escolha do modo de geração
    use_generator = generator is not None

    for step in range(max_new_tokens):
        metrics["steps"] += 1

        # Truncar input para max_seq (janela deslizante simples)
        input_a = input_a[:, -max_seq:]
        input_b = input_b[:, -max_seq:]

        # --- Forward ---
        try:
            if use_generator:
                # GERAÇÃO AUTOREGRESSIVA REAL via GeneratorCNNBiGRU
                # O generator tem decoder GRU + Bahdanau attention
                # Passamos target_proof=None para que ele gere livremente
                # mas limitamos max_len ao tamanho atual + 1
                gen_out = generator(
                    input_a, input_b,
                    target_proof=None,
                    bos_id=bos_id,
                    eos_id=eos_id,
                    max_len=1,  # gera 1 token por vez
                )
                logits = gen_out["logits"]  # [B, 1, V]
                logits = logits[:, -1, :]  # [1, V]
            else:
                # FALLBACK: MultimodalCNNBiGRU em modo "generate" (single-step)
                out = model(
                    input_a, input_b,
                    images=images, audios=audios,
                    mode="generate",
                )
                logits = out["logits"]  # [B, V] ou [B, T, V]
                if logits.dim() == 3:
                    logits = logits[:, -1, :]
            logits = logits.clone()  # [1, V]
        except Exception as e:
            logger.warning(f"Forward falhou no step {step}: {e}")
            break

        # --- Greedy (temperature == 0) ---
        if temperature <= 0.0:
            next_token = torch.argmax(logits, dim=-1, keepdim=True)
            tok_id = next_token.item()
            if tok_id == eos_id:
                break
            generated_ids.append(tok_id)
            token_counts[tok_id] = token_counts.get(tok_id, 0) + 1
            metrics["fallback_to_greedy"] += 1
            # Anexa ao input (trunca para max_seq)
            input_a = torch.cat([input_a, next_token], dim=1)[:, -max_seq:]
            input_b = torch.cat([input_b, next_token], dim=1)[:, -max_seq:]
            continue

        # --- Aplicar filtros de sampling ---
        logits, filter_metrics = apply_sampling_filters(
            logits, token_counts,
            temperature=temperature,
            top_k=top_k,
            top_p=top_p,
            presence_penalty=presence_penalty,
            frequency_penalty=frequency_penalty,
        )
        metrics["tokens_after_topk"] = filter_metrics["tokens_after_topk"]
        metrics["tokens_after_topp"] = filter_metrics["tokens_after_topp"]

        # --- Amostragem Multinomial ---
        # Substituir -inf por um valor muito baixo (evita NaN em softmax)
        logits = torch.where(torch.isinf(logits), torch.full_like(logits, -1e9), logits)
        probs = F.softmax(logits, dim=-1)
        # Lidar com probs NaN (todas -inf)
        if torch.isnan(probs).any() or probs.sum() == 0:
            logger.warning(f"NaN/zero probs no step {step}; usando argmax")
            next_token = torch.argmax(logits, dim=-1, keepdim=True)
        else:
            try:
                next_token = torch.multinomial(probs, num_samples=1)
            except Exception as e:
                logger.warning(f"Multinomial falhou: {e}; usando argmax")
                next_token = torch.argmax(probs, dim=-1, keepdim=True)

        tok_id = next_token.item()
        if tok_id == eos_id:
            break

        generated_ids.append(tok_id)
        token_counts[tok_id] = token_counts.get(tok_id, 0) + 1

        # Anexa ao input (trunca para max_seq)
        input_a = torch.cat([input_a, next_token], dim=1)[:, -max_seq:]
        input_b = torch.cat([input_b, next_token], dim=1)[:, -max_seq:]

    # Decodifica
    text = tokenizer.decode(generated_ids, skip_special=True)

    return {
        "text": text,
        "token_ids": generated_ids,
        "metrics": metrics,
        "used_generator": use_generator,
    }


# ============================================================================
# Avaliação de Perplexidade (V2.0 — CORRIGIDA)
# ============================================================================

@torch.no_grad()
def evaluate_perplexity(
    model: nn.Module,
    dataloader,
    tokenizer,
    device: str = "cpu",
    max_batches: int = 10,
    generator: Optional[GeneratorCNNBiGRU] = None,
) -> Dict[str, float]:
    """Avalia perplexidade do modelo em um dataloader.

    PPL = exp(loss_media)

    VERSÃO 2.0:
      - Se `generator` for fornecido, usa o generator com teacher forcing
        para gerar logits autoregressivos REAIS por timestep.
      - Caso contrário, faz fallback para MultimodalCNNBiGRU (single-step
        broadcast — métrica aproximada).
    """
    model.eval()
    if generator is not None:
        generator.eval()

    total_loss = 0.0
    total_tokens = 0
    pad_id = tokenizer.pad_id
    bos_id = tokenizer.bos_id
    eos_id = tokenizer.eos_id

    n_batches = 0
    for batch_idx, batch in enumerate(dataloader):
        if batch_idx >= max_batches:
            break

        input_ids_a = batch["input_ids_a"].to(device)
        input_ids_b = batch["input_ids_b"].to(device)
        attn_mask_a = batch["attn_mask_a"].to(device)
        attn_mask_b = batch["attn_mask_b"].to(device)

        try:
            if generator is not None:
                # USAR GENERATOR: geração autoregressiva com teacher forcing
                # target_proof = input_ids_b → logits por timestep
                gen_out = generator(
                    input_ids_a, input_ids_b,
                    target_proof=input_ids_b,  # teacher forcing
                    bos_id=bos_id,
                    eos_id=eos_id,
                )
                logits = gen_out["logits"]  # [B, T, V]
            else:
                # FALLBACK: single-step broadcast (NÃO É autoregressivo real)
                # Apenas para compatibilidade
                images = batch.get("images")
                audios = batch.get("audios")
                if images is not None:
                    images = images.to(device)
                if audios is not None:
                    audios = audios.to(device)
                out = model(
                    input_ids_a, input_ids_b,
                    images=images, audios=audios,
                    attn_mask_a=attn_mask_a, attn_mask_b=attn_mask_b,
                    mode="generate",
                )
                # out["logits"] é [B, V] (single-step) → expandir para [B, T, V]
                B = out["logits"].size(0)
                T = input_ids_b.size(1)
                V = out["logits"].size(-1)
                logits = out["logits"].unsqueeze(1).expand(-1, T, -1)

            # Calcular loss por token
            B, T, V = logits.shape
            loss = F.cross_entropy(
                logits.reshape(B * T, V),
                input_ids_b.reshape(B * T),
                ignore_index=pad_id,
                reduction="sum",
            )
            n_tokens = (input_ids_b != pad_id).sum().item()
            total_loss += loss.item()
            total_tokens += n_tokens
            n_batches += 1
        except Exception as e:
            logger.warning(f"PPL batch {batch_idx} falhou: {e}")
            continue

    if total_tokens == 0:
        return {"loss": float("inf"), "ppl": float("inf"), "n_batches": 0}

    avg_loss = total_loss / total_tokens
    try:
        ppl = math.exp(min(avg_loss, 20))
    except OverflowError:
        ppl = float("inf")

    return {
        "loss": avg_loss,
        "ppl": ppl,
        "n_batches": n_batches,
        "used_generator": generator is not None,
    }


# ============================================================================
# Helper: criar context window padrão
# ============================================================================

def make_default_context_window(
    max_window: int = 256,
    n_sink: int = 4,
    embed_dim: int = 64,
    n_heads: int = 4,
    n_layers: int = 2,
    device: str = "cpu",
) -> ContextWindowManager:
    """Factory para criar um ContextWindowManager com defaults sensatos."""
    config = ContextWindowConfig(
        max_window=max_window,
        eviction_strategy="sink_sliding",
        n_sink_tokens=n_sink,
        embed_dim=embed_dim,
        n_heads=n_heads,
        n_layers=n_layers,
        device=device,
    )
    return ContextWindowManager(config)


__all__ = [
    "apply_sampling_filters",
    "generate_with_sampling",
    "evaluate_perplexity",
    "make_default_context_window",
]