Update handler.py

handler.py

@@ -13,24 +13,25 @@ class EndpointHandler():
         self.model = PeftModel.from_pretrained(model, path)
 
     def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
-            """
-            Args:
-                data (Dict): The payload with the text prompt 
-            and generation parameters.
-            """
-            # Get inputs
-            prompt = data.pop("inputs", None)
-            parameters = data.pop("parameters", None)
-            if prompt is None:
-                raise ValueError("Missing prompt.")
-            # Preprocess
-            input_ids = self.tokenizer(prompt, return_tensors="pt").input_ids.to(device)
-            # Forward
-            if parameters is not None:
-                output = self.model.generate(input_ids=input_ids, **parameters)
-            else:
-                output = self.model.generate(input_ids=input_ids)
-            # Postprocess
-            prediction = self.tokenizer.decode(output[0])
-            return {"generated_text": prediction}
+        """
+        Args:
+            data (Dict): The payload with the text prompt 
+        and generation parameters.
+        """
+        # Get inputs
+        prompt = data.pop("inputs", None)
+        parameters = data.pop("parameters", None)
+        if prompt is None:
+            raise ValueError("Missing prompt.")
+        # Preprocess
+        input_ids = self.tokenizer(prompt, return_tensors="pt").input_ids.to(device)
+        # Forward
+        # if parameters is not None:
+        #     output = self.model.generate(input_ids=input_ids, **parameters)
+        # else:
+        #     output = self.model.generate(input_ids=input_ids)
+        output = self.model.generate(input_ids, temperature=0.9, max_new_tokens=50)
+        # Postprocess
+        prediction = self.tokenizer.batch_decode(outputs.detach().cpu().numpy(), skip_special_tokens=True)[0]
+        return {"generated_text": prediction}