Update handler.py

handler.py

@@ -3,11 +3,15 @@ from transformers import pipeline, AutoTokenizer, BartForConditionalGeneration
 
 class EndpointHandler():
     def __init__(self, path=""):
-        self.model = BartForConditionalGeneration.from_pretrained(path)
-        self.tokenizer = AutoTokenizer(path)
+        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        try:
+            self.model = BartForConditionalGeneration.from_pretrained(path).to(self.device)
+            self.tokenizer = AutoTokenizer.from_pretrained(path)
+        except Exception as e:
+            print(f"Error loading model or tokenizer from path {path}: {e}")
         
 
-    def __call__(self, data: str) -> str:
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         """
        data args:
             inputs (:obj: `str`)
@@ -18,11 +22,15 @@ class EndpointHandler():
         # get inputs
   
 
-        text_tokenized = self.tokenizer(
-        [data], padding="max_length", truncation=True, max_length=1024,return_tensors='pt')
-    
-        prediction_token = self.model.generate(text_tokenized["input_ids"], max_length = 256, num_beams = 6)
-        
-        prediction_summary = self.tokenizer.decode(prediction_token[0][2:-1:1])
+        inputs = data.get("inputs", "")
+        if not inputs:
+            return [{"error": "No inputs provided"}]
+
+        tokenized_input = self.tokenizer(inputs, return_tensors="pt", truncation=True, max_length=1024, padding="max_length")
+        tokenized_input = tokenized_input.to(self.device)  # Move input tensors to the same device as model
+
+        summary_ids = self.model.generate(**tokenized_input, max_length=256, do_sample=True, top_p=0.8)
+
+        summary_text = self.tokenizer.decode(summary_ids[0], skip_special_tokens=True)
 
-        return prediction_summary
+        return [{"summary": summary_text}]