handler.py

from transformers import pipeline, AutoTokenizer, AutoModelForQuestionAnswering

class EndpointHandler:
    def __init__(self, path=""):
        self.model = AutoModelForQuestionAnswering.from_pretrained(path)
        self.tokenizer = AutoTokenizer.from_pretrained(path)
        self.qa_pipeline = pipeline("question-answering", model=self.model, tokenizer=self.tokenizer)

    def __call__(self, data):
        inputs = data.get("inputs", {})
        question = inputs.get("question")
        context = inputs.get("context")
        
        if not question or not context:
            return {"error": "Both 'question' and 'context' must be provided."}
        
        result = self.qa_pipeline(question=question, context=context)
        return result