Update handler

handler.py

@@ -24,16 +24,18 @@ class EndpointHandler():
         """
         LOGGER.info(f"Received data: {data}")
         # Get inputs
-        prompt = data.pop("prompt", data)
+        prompt = data.pop("prompt", None)
         parameters = data.pop("parameters", None)
+        if prompt is None:
+            raise ValueError("Missing prompt.")
         # Preprocess
-        input = self.tokenizer(prompt, return_tensors="pt")
+        inputs = self.tokenizer(prompt, return_tensors="pt")
         # Forward
         LOGGER.info(f"Start generation.")
         if parameters is not None:
-            output = self.model.generate(**input, **parameters)
+            output = self.model.generate(**inputs, **parameters)
         else:
-            output = self.model.generate(**input)
+            output = self.model.generate(**inputs)
         # Postprocess
         prediction = self.tokenizer.decode(output[0])
         LOGGER.info(f"Generated text: {prediction}")