Updates handler.py function names to match requirements

handler.py

@@ -1,59 +1,38 @@
-import torch
 from transformers import AutoTokenizer, AutoModelForQuestionAnswering
-from typing import Dict, Any
+import torch
 
-class QuestionAnsweringHandler:
-    def __init__(self):
-        self.model = None
-        self.tokenizer = None
+class EndpointHandler:
+    def __init__(self, model_path=""):
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
+        self.model = AutoModelForQuestionAnswering.from_pretrained(model_path).to(self.device)
 
-    def initialize(self, ctx):
-        model_dir = ctx.system_properties.get("model_dir")
-        self.model = AutoModelForQuestionAnswering.from_pretrained(model_dir).to(self.device)
-        self.model.eval()
-        self.tokenizer = AutoTokenizer.from_pretrained(model_dir)
-
-    def preprocess(self, data: Any) -> Dict[str, str]:
-        # Expect JSON with {"question": ..., "context": ...}
-        question = data[0]["body"].get("question", "")
-        context = data[0]["body"].get("context", "")
-        return {"question": question, "context": context}
-
-    def inference(self, inputs: Dict[str, str]) -> Dict[str, str]:
-        question = inputs["question"]
-        context = inputs["context"]
-
-        encoded = self.tokenizer(
-            question,
-            context,
-            return_tensors="pt",
-            max_length=512,
-            truncation=True
-        ).to(self.device)
+    def __call__(self, data):
+        """
+        data: dict containing 'inputs' with 'question' and 'context' keys
+        """
+        inputs = data.get("inputs", {})
+        question = inputs.get("question")
+        context = inputs.get("context")
 
-        with torch.no_grad():
-            outputs = self.model(**encoded)
+        if not question or not context:
+            return {"error": "Missing question or context"}
 
-        start_logits = outputs.start_logits[0]
-        end_logits = outputs.end_logits[0]
+        encoded = self.tokenizer(question, context, return_tensors="pt", truncation=True, max_length=512).to(self.device)
+
+        with torch.no_grad():
+            output = self.model(**encoded)
 
-        max_answer_len = 30
-        input_ids = encoded["input_ids"][0]
+        start_scores = output.start_logits[0]
+        end_scores = output.end_logits[0]
 
-        # Score spans and find the best one
-        best_score = float("-inf")
-        best_span = ""
+        # Get best span
+        start_idx = torch.argmax(start_scores)
+        end_idx = torch.argmax(end_scores)
 
-        for start in range(len(start_logits)):
-            for end in range(start, min(start + max_answer_len, len(end_logits))):
-                score = start_logits[start] + end_logits[end]
-                if score > best_score:
-                    best_score = score
-                    span_ids = input_ids[start:end + 1]
-                    best_span = self.tokenizer.decode(span_ids, skip_special_tokens=True)
+        if end_idx < start_idx:
+            return {"answer": ""}
 
-        return {"best_span": best_span.strip()}
+        answer = self.tokenizer.decode(encoded["input_ids"][0][start_idx:end_idx + 1], skip_special_tokens=True)
 
-    def postprocess(self, output: Dict[str, str]) -> [Dict[str, str]]:
-        return [output]
+        return {"answer": answer}