Create model_utils.py

model_utils.py

@@ -0,0 +1,60 @@
+# model_utils.py
+import torch
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+
+# 🔁 Mets ici le chemin ou le repo HF de ton modèle
+MODEL_NAME = "ton-username/stackoverflow-tags-bert"
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+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)}
+
+
+def _to_device(batch):
+    return {k: v.to(device) for k, v in batch.items()}
+
+
+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.
+    """
+    enc = tokenizer(
+        text,
+        return_tensors="pt",
+        truncation=True,
+        padding=True,
+        max_length=256,
+    )
+
+    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 = 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)}"),
+            "score": float(probs[i]),
+        }
+        for i in indices
+    ]
+    return results