Update handler

handler.py

@@ -23,15 +23,15 @@ class EndpointHandler():
         """
         LOGGER.info(f"Received data: {data}")
         # Get inputs
-        prompt = data.pop("prompt", data)
+        inputs = data.pop("inputs", data)
         parameters = data.pop("parameters", None)
         # Preprocess
-        inputs = self.tokenizer(prompt, return_tensors="pt")
+        inputs_ids = self.tokenizer(inputs, return_tensors="pt").inputs_ids
         # Forward
         if parameters is not None:
-            outputs = self.model.generate(**inputs, **parameters)
+            outputs = self.model.generate(inputs_ids, **parameters)
         else:
-            outputs = self.model.generate(**inputs)
+            outputs = self.model.generate(inputs_ids)
         # Postprocess
         prediction = self.tokenizer.decode(outputs[0])
         LOGGER.info(f"Generated text: {prediction}")