Update model_utils.py

model_utils.py

@@ -1,8 +1,8 @@
 # model_utils.py
 import torch
+import joblib
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
-# 🔁 Mets ici le chemin ou le repo HF de ton modèle
 MODEL_NAME = "maxcasado/BERT_overflow"
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -12,10 +12,9 @@ model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
 model.to(device)
 model.eval()
 
-# On essaie de récupérer les labels à partir de la config
-id2label = getattr(model.config, "id2label", None)
-if not id2label:
-    id2label = {i: f"label_{i}" for i in range(model.config.num_labels)}
+# MultiLabelBinarizer pour récupérer les noms de tags
+mlb = joblib.load("mlb.joblib")  # le fichier présent dans ton repo modèle
+classes = list(mlb.classes_)     # index -> nom de tag
 
 
 def _to_device(batch):
@@ -24,8 +23,7 @@ def _to_device(batch):
 
 def predict_proba(text: str, top_k: int = 10):
     """
-    Prend une question en entrée, renvoie les top_k tags avec leurs probas.
-    Gère multi-class et multi-label.
+    Multi-label : renvoie top_k tags avec proba (sigmoid).
     """
     enc = tokenizer(
         text,
@@ -38,21 +36,15 @@ def predict_proba(text: str, top_k: int = 10):
     with torch.no_grad():
         outputs = model(**_to_device(enc))
         logits = outputs.logits[0]
-
-        # Heuristique : si problème multi-label
-        if getattr(model.config, "problem_type", None) == "multi_label_classification":
-            probs = torch.sigmoid(logits)
-        else:
-            probs = torch.softmax(logits, dim=-1)
+        probs = torch.sigmoid(logits)  # multi-label !
 
     probs = probs.cpu().numpy()
 
-    # indices triés par proba décroissante
     indices = probs.argsort()[::-1][:top_k]
 
     results = [
         {
-            "label": id2label.get(int(i), f"label_{int(i)}"),
+            "label": classes[int(i)],
             "score": float(probs[i]),
         }
         for i in indices