refactor: single-fold fold_04 + Platt scaling (remove ensemble)

inference.py

@@ -1,13 +1,13 @@
-"""Carregamento do modelo e inferência (ensemble multi-fold calibrado).
+"""Carregamento do modelo e inferência (bge-m3 FT-Solo, single-fold calibrado).
 
-Usa o mecanismo nativo do PEFT (load_adapter / set_adapter) para trocar
-adapters em um único encoder, evitando sobreposição de pesos entre folds.
+Platt scaling pós-treino: P_calib = sigmoid(CALIB_A * logit(P_raw) + CALIB_B).
+Com CALIB_A=1.0, CALIB_B=0.0 (defaults) a transformação é identidade.
 """
 from __future__ import annotations
 
 import logging
 from functools import lru_cache
-from typing import Iterable, List, Tuple
+from typing import Iterable
 
 import numpy as np
 import torch
@@ -17,22 +17,19 @@ from peft import PeftModel
 from transformers import AutoModel, AutoTokenizer
 
 from config import (
-    ADAPTER_DIRNAME,
-    ARTIFACTS_DIR,
+    ADAPTER_PATH,
     BATCH_SIZE,
     CALIB_A,
     CALIB_B,
-    HEAD_FILENAME,
+    HEAD_PATH,
     HF_TOKEN,
     MAX_LENGTH,
-    MODEL_FOLDS,
     MODEL_NAME,
-    TEMPERATURE,
 )
 
 logger = logging.getLogger(__name__)
 
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+DEVICE    = "cuda" if torch.cuda.is_available() else "cpu"
 AMP_DTYPE = (
     (torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16)
     if DEVICE == "cuda"
@@ -52,15 +49,13 @@ def mean_pool(last_hidden_states: torch.Tensor, attention_mask: torch.Tensor) ->
 
 
 @lru_cache(maxsize=1)
-def load_models() -> Tuple[AutoTokenizer, PeftModel, List[nn.Module]]:
-    """
-    Carrega tokenizer, encoder com TODOS os adapters (via PEFT multi-adapter),
-    e lista de cabeças lineares, uma por fold.
-
-    O encoder base é carregado UMA VEZ. Cada adapter é adicionado com
-    `load_adapter(adapter_name=fold)` — sem sobreposição de pesos.
-    A troca de adapter em inferência usa `encoder.set_adapter(fold)`.
-    """
+def load_model():
+    """Retorna (tokenizer, encoder, head). Carregado uma única vez por processo."""
+    if not ADAPTER_PATH.exists():
+        raise FileNotFoundError(f"Adapter LoRA não encontrado em {ADAPTER_PATH}.")
+    if not HEAD_PATH.exists():
+        raise FileNotFoundError(f"Cabeça classificadora não encontrada em {HEAD_PATH}.")
+
     logger.info("Carregando tokenizer de %s", MODEL_NAME)
     tokenizer = AutoTokenizer.from_pretrained(
         MODEL_NAME, padding_side="right", token=HF_TOKEN
@@ -68,117 +63,90 @@ def load_models() -> Tuple[AutoTokenizer, PeftModel, List[nn.Module]]:
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
 
-    logger.info("Carregando encoder base %s (dtype=%s)", MODEL_NAME, AMP_DTYPE)
+    logger.info("Carregando encoder base %s (dtype=%s, device=%s)", MODEL_NAME, AMP_DTYPE, DEVICE)
     base_encoder = AutoModel.from_pretrained(
         MODEL_NAME, low_cpu_mem_usage=True, torch_dtype=AMP_DTYPE, token=HF_TOKEN
     ).to(DEVICE)
 
-    # Carrega o PRIMEIRO adapter — cria o PeftModel
-    first_fold = MODEL_FOLDS[0]
-    first_adapter_dir = ARTIFACTS_DIR / ADAPTER_DIRNAME.format(fold=first_fold)
-    assert first_adapter_dir.exists(), f"{first_adapter_dir} não encontrado"
+    logger.info("Anexando adapter LoRA de %s", ADAPTER_PATH)
     encoder = PeftModel.from_pretrained(
-        base_encoder, str(first_adapter_dir),
-        adapter_name=first_fold, is_trainable=False,
+        base_encoder, str(ADAPTER_PATH), is_trainable=False
     ).to(DEVICE)
-
-    # Carrega os DEMAIS adapters no mesmo PeftModel — sem modificar o base
-    for fold in MODEL_FOLDS[1:]:
-        adapter_dir = ARTIFACTS_DIR / ADAPTER_DIRNAME.format(fold=fold)
-        assert adapter_dir.exists(), f"{adapter_dir} não encontrado"
-        encoder.load_adapter(str(adapter_dir), adapter_name=fold)
-
     encoder.eval()
-    logger.info("%d adapters carregados: %s", len(MODEL_FOLDS), MODEL_FOLDS)
 
-    # Carrega as cabeças lineares
-    heads: List[nn.Module] = []
-    for fold in MODEL_FOLDS:
-        head_path = ARTIFACTS_DIR / HEAD_FILENAME.format(fold=fold)
-        assert head_path.exists(), f"{head_path} não encontrada"
-        payload = torch.load(head_path, map_location="cpu")
-        state = payload.get("state_dict", payload) if isinstance(payload, dict) else payload
-        head = nn.Linear(int(state["weight"].shape[1]), 1)
-        head.load_state_dict(state)
-        heads.append(head.to(DEVICE).eval())
+    logger.info("Carregando cabeça linear de %s", HEAD_PATH)
+    payload    = torch.load(HEAD_PATH, map_location="cpu")
+    head_state = payload.get("state_dict", payload) if isinstance(payload, dict) else payload
+    in_feat    = int(head_state["weight"].shape[1])
+    head       = nn.Linear(in_feat, 1)
+    head.load_state_dict(head_state)
+    head = head.to(DEVICE).eval()
 
-    return tokenizer, encoder, heads
+    logger.info("Modelo pronto. In_features da cabeça: %d", in_feat)
+    return tokenizer, encoder, head
 
 
 def warmup() -> None:
-    """Força carregamento imediato para evitar cold-start no primeiro request."""
-    load_models()
+    """Força carregamento imediato para evitar cold-start."""
+    load_model()
 
 
 @torch.no_grad()
 def predict_batch(texts: Iterable[str], batch_size: int = BATCH_SIZE) -> np.ndarray:
-    """
-    Probabilidade calibrada de 'útil', em média entre todos os folds.
+    """Probabilidade calibrada de 'útil' para cada texto. Shape (N,)."""
+    tokenizer, encoder, head = load_model()
 
-    Para cada fold: ativa o adapter com set_adapter(), roda o forward,
-    aplica Platt scaling se configurado, acumula. Depois tira a média.
-    """
     if isinstance(texts, str):
         texts = [texts]
     texts = list(texts)
     if not texts:
         return np.zeros(0, dtype=np.float64)
 
-    tokenizer, encoder, heads = load_models()
+    preds = []
     autocast_device = "cuda" if DEVICE == "cuda" else "cpu"
-    fold_preds: List[np.ndarray] = []
-
-    for fold, head in zip(MODEL_FOLDS, heads):
-        # Ativa SOMENTE o adapter deste fold; os demais ficam inativos
-        encoder.set_adapter(fold)
-
-        preds = []
-        for i in range(0, len(texts), batch_size):
-            batch = texts[i: i + batch_size]
-            instr = [build_instruction_text(t) for t in batch]
-            toks = tokenizer(
-                instr, padding=True, truncation=True,
-                max_length=MAX_LENGTH, return_tensors="pt",
-            ).to(DEVICE)
-
-            with torch.inference_mode(), torch.autocast(
-                device_type=autocast_device, dtype=AMP_DTYPE, enabled=(DEVICE == "cuda")
-            ):
-                out = encoder(**toks)
-                emb = mean_pool(out.last_hidden_state, toks["attention_mask"])
-                emb = F.normalize(emb, p=2, dim=1)
-                logits = head(emb.to(head.weight.dtype)).squeeze(-1)
-                if TEMPERATURE != 1.0:
-                    logits = logits / TEMPERATURE
-                p = torch.sigmoid(logits).float().cpu().numpy()
-            preds.append(p)
-
-        p_fold = np.clip(np.concatenate(preds).astype(np.float64), 1e-6, 1 - 1e-6)
-
-        # Platt scaling: P_calib = sigmoid(A * logit(p) + B)
-        # NOTA: sigmoid(x) = 1/(1+exp(-x)) — o sinal negativo é obrigatório
-        if CALIB_A != 1.0 or CALIB_B != 0.0:
-            logits_np = np.log(p_fold / (1.0 - p_fold))
-            p_fold = 1.0 / (1.0 + np.exp(-(CALIB_A * logits_np + CALIB_B)))
-
-        fold_preds.append(p_fold)
-
-    return np.mean(fold_preds, axis=0) if len(fold_preds) > 1 else fold_preds[0]
+
+    for i in range(0, len(texts), batch_size):
+        batch = texts[i : i + batch_size]
+        instr = [build_instruction_text(t) for t in batch]
+        toks  = tokenizer(
+            instr, padding=True, truncation=True,
+            max_length=MAX_LENGTH, return_tensors="pt",
+        ).to(DEVICE)
+
+        with torch.inference_mode(), torch.autocast(
+            device_type=autocast_device, dtype=AMP_DTYPE, enabled=(DEVICE == "cuda")
+        ):
+            out  = encoder(**toks)
+            emb  = mean_pool(out.last_hidden_state, toks["attention_mask"])
+            emb  = F.normalize(emb, p=2, dim=1)
+            # Em CPU sem autocast, encoder fp16 + head fp32 → cast necessário
+            logits = head(emb.to(head.weight.dtype)).squeeze(-1)
+            p = torch.sigmoid(logits).float().cpu().numpy()
+        preds.append(p)
+
+    p_raw = np.clip(np.concatenate(preds).astype(np.float64), 1e-6, 1 - 1e-6)
+
+    # Platt scaling: P_calib = sigmoid(A * logit(P_raw) + B)
+    # sigmoid(x) = 1/(1+exp(-x))  — sinal negativo obrigatório no exp
+    if CALIB_A != 1.0 or CALIB_B != 0.0:
+        logit_raw = np.log(p_raw / (1.0 - p_raw))
+        return 1.0 / (1.0 + np.exp(-(CALIB_A * logit_raw + CALIB_B)))
+    return p_raw
 
 
 def predict_one(text: str) -> float:
-    """Probabilidade calibrada para um único texto."""
+    """Atalho: probabilidade calibrada para um único texto."""
     return float(predict_batch([text])[0])
 
 
 def explain_occlusion(text: str, batch_size: int = BATCH_SIZE) -> dict:
-    """Leave-one-out usando o ensemble completo com calibração."""
+    """Leave-one-out por palavra. Δ = P(texto) − P(texto sem a palavra)."""
     words = text.split()
     if not words:
         p = predict_one(text)
         return {"proba_full": p, "tokens": [], "contributions": []}
-    variants = [" ".join(words[:i] + words[i + 1:]) for i in range(len(words))]
-    probs = predict_batch([text] + variants, batch_size=batch_size)
-    p_full = float(probs[0])
+    variants  = [" ".join(words[:i] + words[i + 1:]) for i in range(len(words))]
+    probs     = predict_batch([text] + variants, batch_size=batch_size)
+    p_full    = float(probs[0])
     return {"proba_full": p_full, "tokens": words,
             "contributions": (p_full - probs[1:]).tolist()}