Update handler.py

handler.py

@@ -1,15 +1,18 @@
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
+from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
 import torch
+import os
 
 class EndpointHandler:
-    def __init__(self):
-        self.model = AutoModelForSequenceClassification.from_pretrained(".")
-        self.tokenizer = AutoTokenizer.from_pretrained(".")
+    def __init__(self, model_dir):
+        # Hugging Face passe le répertoire du modèle ici
+        self.model_dir = model_dir
+        self.model = AutoModelForSequenceClassification.from_pretrained(model_dir)
+        self.tokenizer = AutoTokenizer.from_pretrained(model_dir)
         self.labels = ["presentation","projects","skills","education","contact","fallback"]
+        # Optionnel : créer un pipeline pour simplifier l'inférence
+        self.classifier = pipeline("text-classification", model=self.model, tokenizer=self.tokenizer)
 
     def inference(self, inputs):
-        encoded = self.tokenizer(inputs, return_tensors="pt")
-        with torch.no_grad():
-            logits = self.model(**encoded).logits
-        predicted_class = torch.argmax(logits, dim=1).item()
-        return {"label": self.labels[predicted_class], "score": 1.0}
+        # inputs est une chaîne de texte
+        outputs = self.classifier(inputs)
+        return {"label": outputs[0]["label"], "score": float(outputs[0]["score"])}