returning generated ids

handler.py

@@ -38,7 +38,7 @@ class EndpointHandler():
         )
 
         generated_text = self.tokenizer.decode(generated_ids[0][input_ids.shape[1]:], skip_special_tokens=True)
-        prediction = [{"generated_text": generated_text}]
+        prediction = [{"generated_text": generated_text, "generated_ids": generated_ids[0][input_ids.shape[1]:].tolist()}]
         return prediction
 
 

sample.py

@@ -7,7 +7,7 @@ my_handler = EndpointHandler(path=".")
 payload = {"inputs": "I want to turn the next page of the"}
 
 # test the handler
-payload=my_handler(payload, additional_bad_words_ids=[[1234, 2345]])
+payload=my_handler(payload)
 
 # show results
 print("output:", payload)