from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline

class EndpointHandler:
    def __init__(self, model_dir):
        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"]
        self.classifier = pipeline("text-classification", model=self.model, tokenizer=self.tokenizer)

    def __call__(self, request):
        """
        Hugging Face Default container attend __call__ comme point d'entrée.
        `request` est le payload JSON reçu par l'endpoint.
        """
        text = request.get("inputs", "")
        outputs = self.classifier(text)
        return {"label": outputs[0]["label"], "score": float(outputs[0]["score"])}