Upload cnn_bigru/inference/inference.py with huggingface_hub

cnn_bigru/inference/inference.py

@@ -0,0 +1,257 @@
+"""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
+"""
+from __future__ import annotations
+
+import logging
+import math
+from typing import Dict, List, Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+logger = logging.getLogger(__name__)
+
+
+@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",
+) -> Dict[str, any]:
+    """Gera texto com amostragem avançada.
+
+    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
+
+    Returns:
+        dict com:
+            text: texto gerado
+            token_ids: lista de IDs gerados
+            metrics: dicionário de métricas
+    """
+    model.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,
+    }
+
+    for step in range(max_new_tokens):
+        metrics["steps"] += 1
+
+        # Forward
+        try:
+            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, :]  # pega último token
+            logits = logits.clone()  # [1, V]
+        except Exception as e:
+            logger.warning("Forward falhou no step %d: %s", step, e)
+            break
+
+        # Penalidades de repetição
+        if presence_penalty != 0.0 or frequency_penalty != 0.0:
+            for tok_id, count in token_counts.items():
+                logits[0, tok_id] -= (presence_penalty + count * frequency_penalty)
+
+        # Temperatura
+        if temperature > 0.0:
+            logits = logits / temperature
+        else:
+            # Greedy
+            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
+            input_a = torch.cat([input_a, next_token], dim=1)
+            input_b = torch.cat([input_b, next_token], dim=1)
+            continue
+
+        # Top-K
+        if top_k > 0:
+            k = min(top_k, logits.size(-1))
+            top_vals, _ = torch.topk(logits, k)
+            thresh = top_vals[:, -1:]
+            logits[logits < thresh] = float("-inf")
+            metrics["tokens_after_topk"] = k
+
+        # Top-P (Nucleus)
+        if top_p < 1.0:
+            sorted_logits, sorted_idx = torch.sort(logits, descending=True)
+            sorted_probs = F.softmax(sorted_logits, dim=-1)
+            cum_probs = torch.cumsum(sorted_probs, dim=-1)
+
+            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())
+
+        # Amostragem
+        probs = F.softmax(logits, dim=-1)
+        try:
+            next_token = torch.multinomial(probs, num_samples=1)
+        except Exception as e:
+            logger.warning("Multinomial falhou: %s — usando argmax", e)
+            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 (mantém janela)
+        max_seq = 64  # janela máxima
+        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,
+    }
+
+
+@torch.no_grad()
+def evaluate_perplexity(
+    model: nn.Module,
+    dataloader,
+    tokenizer,
+    device: str = "cpu",
+    max_batches: int = 10,
+) -> Dict[str, float]:
+    """Avalia perplexidade do modelo em um dataloader.
+
+    PPL = exp(loss_media)
+    """
+    model.eval()
+    total_loss = 0.0
+    total_tokens = 0
+    pad_id = tokenizer.pad_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)
+        images = batch["images"].to(device)
+        audios = batch["audios"].to(device)
+        attn_mask_a = batch["attn_mask_a"].to(device)
+        attn_mask_b = batch["attn_mask_b"].to(device)
+
+        try:
+            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",
+            )
+            logits = out["logits"]  # [B, V]
+            # Para PPL, usamos input_ids_b como target (próximo token)
+            # Expandimos logits para [B, T, V]
+            B = logits.size(0)
+            T = input_ids_b.size(1)
+            V = logits.size(-1)
+            expanded = logits.unsqueeze(1).expand(-1, T, -1)
+
+            loss = F.cross_entropy(
+                expanded.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("PPL batch %d falhou: %s", batch_idx, 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}
+
+
+__all__ = ["generate_with_sampling", "evaluate_perplexity"]