Update model_utils.py

model_utils.py

@@ -2,19 +2,22 @@
 import torch
 import joblib
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
+from huggingface_hub import hf_hub_download
 
 MODEL_NAME = "maxcasado/BERT_overflow"
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
+print("Loading tokenizer and model...")
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
 model.to(device)
 model.eval()
 
-# 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
+print("Loading MultiLabelBinarizer (mlb.joblib)...")
+mlb_path = hf_hub_download(MODEL_NAME, "mlb.joblib")
+mlb = joblib.load(mlb_path)
+classes = list(mlb.classes_)  # index -> tag name
 
 
 def _to_device(batch):
@@ -23,7 +26,9 @@ def _to_device(batch):
 
 def predict_proba(text: str, top_k: int = 10):
     """
-    Multi-label : renvoie top_k tags avec proba (sigmoid).
+    Multi-label prediction:
+    - entrée : texte de la question
+    - sortie : top_k tags avec leurs probabilités (sigmoid)
     """
     enc = tokenizer(
         text,
@@ -35,18 +40,16 @@ def predict_proba(text: str, top_k: int = 10):
 
     with torch.no_grad():
         outputs = model(**_to_device(enc))
-        logits = outputs.logits[0]
-        probs = torch.sigmoid(logits)  # multi-label !
+        logits = outputs.logits[0]           # shape [num_labels]
+        probs = torch.sigmoid(logits)        # multi-label
 
     probs = probs.cpu().numpy()
-
     indices = probs.argsort()[::-1][:top_k]
 
-    results = [
+    return [
         {
             "label": classes[int(i)],
             "score": float(probs[i]),
         }
         for i in indices
     ]
-    return results