"""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", ]