fix: ensemble via set_adapter (elimina sobreposição de adapters) + Platt scaling com sinal correto

inference.py

@@ -1,10 +1,7 @@
-"""Carregamento do modelo e inferência (calibrado e ensemblado).
+"""Carregamento do modelo e inferência (ensemble multi-fold calibrado).
 
-Serve o ensemble calibrado com base BAAI/bge-m3 + LoRA. 
-Carrega todas as combinações definidas em config.MODEL_FOLDS, roda a
-inferência para cada uma e calcula a média das probabilidades.
-As probabilidades brutas passam por uma transformação paramétrica 
-(Platt scaling / temperature scaling).
+Usa o mecanismo nativo do PEFT (load_adapter / set_adapter) para trocar
+adapters em um único encoder, evitando sobreposição de pesos entre folds.
 """
 from __future__ import annotations
 
@@ -20,34 +17,31 @@ from peft import PeftModel
 from transformers import AutoModel, AutoTokenizer
 
 from config import (
+    ADAPTER_DIRNAME,
     ARTIFACTS_DIR,
     BATCH_SIZE,
     CALIB_A,
     CALIB_B,
     HEAD_FILENAME,
-    MODEL_FOLDS,
-    ADAPTER_DIRNAME,
     HF_TOKEN,
     MAX_LENGTH,
+    MODEL_FOLDS,
     MODEL_NAME,
     TEMPERATURE,
 )
 
 logger = logging.getLogger(__name__)
 
-# ---------------------------------------------------------------------------
-# Dispositivo e dtype
-# ---------------------------------------------------------------------------
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-
-if DEVICE == "cuda":
-    AMP_DTYPE = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
-else:
-    AMP_DTYPE = torch.float16
+AMP_DTYPE = (
+    (torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16)
+    if DEVICE == "cuda"
+    else torch.float16
+)
 
 
 def build_instruction_text(text: str) -> str:
-    """Retorna o texto sem prompt de instrução (bge-m3 não usa prompts)."""
+    """bge-m3 não usa prompt de instrução — retorna o texto cru."""
     return text if isinstance(text, str) else ""
 
 
@@ -58,86 +52,95 @@ def mean_pool(last_hidden_states: torch.Tensor, attention_mask: torch.Tensor) ->
 
 
 @lru_cache(maxsize=1)
-def load_models() -> List[Tuple[AutoTokenizer, PeftModel, nn.Module]]:
-    """
-    Carrega todas as combinações (tokenizer, encoder, head) definidas
-    em config.MODEL_FOLDS. Retorna uma lista de tuplas. 
-    O tokenizer e o encoder base são compartilhados para economizar memória.
+def load_models() -> Tuple[AutoTokenizer, PeftModel, List[nn.Module]]:
     """
-    models = []
+    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)`.
+    """
     logger.info("Carregando tokenizer de %s", MODEL_NAME)
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, padding_side="right", token=HF_TOKEN)
+    tokenizer = AutoTokenizer.from_pretrained(
+        MODEL_NAME, padding_side="right", token=HF_TOKEN
+    )
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
 
-    logger.info(
-        "Carregando encoder base %s (dtype=%s, device=%s)", MODEL_NAME, AMP_DTYPE, DEVICE
-    )
+    logger.info("Carregando encoder base %s (dtype=%s)", MODEL_NAME, AMP_DTYPE)
     base_encoder = AutoModel.from_pretrained(
-        MODEL_NAME,
-        low_cpu_mem_usage=True,
-        torch_dtype=AMP_DTYPE,
-        token=HF_TOKEN,
+        MODEL_NAME, low_cpu_mem_usage=True, torch_dtype=AMP_DTYPE, token=HF_TOKEN
     ).to(DEVICE)
 
-    for fold in MODEL_FOLDS:
+    # 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"
+    encoder = PeftModel.from_pretrained(
+        base_encoder, str(first_adapter_dir),
+        adapter_name=first_fold, 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)
-        head_path = ARTIFACTS_DIR / HEAD_FILENAME.format(fold=fold)
-        if not adapter_dir.exists() or not head_path.exists():
-            raise FileNotFoundError(
-                f"Artifacts do fold '{fold}' não encontrados em {adapter_dir} e {head_path}"
-            )
-        logger.info("Anexando adapter LoRA de %s", adapter_dir)
-        encoder = PeftModel.from_pretrained(base_encoder, str(adapter_dir), is_trainable=False).to(DEVICE)
-        encoder.eval()
+        assert adapter_dir.exists(), f"{adapter_dir} não encontrado"
+        encoder.load_adapter(str(adapter_dir), adapter_name=fold)
 
-        logger.info("Carregando cabeça linear de %s", head_path)
-        head_payload = torch.load(head_path, map_location="cpu")
-        head_state = head_payload.get("state_dict", head_payload) if isinstance(head_payload, dict) else head_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()
+    encoder.eval()
+    logger.info("%d adapters carregados: %s", len(MODEL_FOLDS), MODEL_FOLDS)
 
-        models.append((tokenizer, encoder, head))
+    # 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("%d modelos de ensemble carregados.", len(models))
-    return models
+    return tokenizer, encoder, heads
 
 
 def warmup() -> None:
-    """Carrega todos os modelos imediatamente para evitar cold-start."""
+    """Força carregamento imediato para evitar cold-start no primeiro request."""
     load_models()
 
 
 @torch.no_grad()
 def predict_batch(texts: Iterable[str], batch_size: int = BATCH_SIZE) -> np.ndarray:
-    """Retorna a probabilidade calibrada de 'útil' para cada texto, em média entre folds."""
+    """
+    Probabilidade calibrada de 'útil', em média entre todos os folds.
+
+    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)
 
-    # Coleção de predições por fold
-    fold_preds: List[np.ndarray] = []
-    models = load_models()
-    # Determina dtype para autocast
+    tokenizer, encoder, heads = load_models()
     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)
 
-    for tokenizer, encoder, head in models:
         preds = []
         for i in range(0, len(texts), batch_size):
-            batch = texts[i : i + 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",
+                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")
             ):
@@ -145,48 +148,37 @@ def predict_batch(texts: Iterable[str], batch_size: int = BATCH_SIZE) -> np.ndar
                 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)
-                # Temperature scaling (divide os logits por TEMPERATURE)
                 if TEMPERATURE != 1.0:
                     logits = logits / TEMPERATURE
-                # Calcula p via sigmóide nos logits (pré-calibração)
                 p = torch.sigmoid(logits).float().cpu().numpy()
             preds.append(p)
-        preds_full = np.concatenate(preds).astype(np.float64)
-        # Clip para evitar 0 ou 1 exatos
-        preds_full = np.clip(preds_full, 1e-6, 1 - 1e-6)
-        # Converte p em logit para aplicar calibração Platt: z = log(p/(1-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(preds_full / (1.0 - preds_full))
-            calibrated = 1.0 / (1.0 + np.exp(CALIB_A * logits_np + CALIB_B))
-        else:
-            calibrated = preds_full
-        fold_preds.append(calibrated)
+            logits_np = np.log(p_fold / (1.0 - p_fold))
+            p_fold = 1.0 / (1.0 + np.exp(-(CALIB_A * logits_np + CALIB_B)))
 
-    # Média do Ensemble
-    if len(fold_preds) > 1:
-        final = np.mean(fold_preds, axis=0)
-    else:
-        final = fold_preds[0]
-    return final
+        fold_preds.append(p_fold)
+
+    return np.mean(fold_preds, axis=0) if len(fold_preds) > 1 else fold_preds[0]
 
 
 def predict_one(text: str) -> float:
-    """Retorna a probabilidade calibrada para um único texto."""
+    """Probabilidade calibrada para um único texto."""
     return float(predict_batch([text])[0])
 
 
 def explain_occlusion(text: str, batch_size: int = BATCH_SIZE) -> dict:
-    """
-    Explicação leave-one-out por palavra, usando a média do ensemble e aplicando calibração.
-    Δ = P(texto completo) − P(texto sem a palavra).
-    """
+    """Leave-one-out usando o ensemble completo com calibração."""
     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))]
-    all_texts = [text] + variants
-    probs = predict_batch(all_texts, batch_size=batch_size)
+    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])
-    contribs = (p_full - probs[1:]).tolist()
-    return {"proba_full": p_full, "tokens": words, "contributions": contribs}
+    return {"proba_full": p_full, "tokens": words,
+            "contributions": (p_full - probs[1:]).tolist()}