piercemaloney
/

llemma_7b

Text Generation

text-generation-inference

Inference Endpoints

Model card Files Files and versions Community

Pierce Maloney commited on Apr 10

Commit

88fdb99

•

1 Parent(s): 2872c1c

returning ids

Files changed (1) hide show

handler.py +1 -2

handler.py CHANGED Viewed

@@ -37,9 +37,8 @@ class EndpointHandler():
         # Decode the generated ids to text
         # Exclude the input_ids length to get only the new tokens
-        print("Generated IDs:", prediction_ids[0, input_ids.shape[1]:])
         prediction_text = self.tokenizer.decode(prediction_ids[0, input_ids.shape[1]:], skip_special_tokens=True)
-        return [{"generated_text": prediction_text}]
 class StopAtPeriodCriteria(StoppingCriteria):

         # Decode the generated ids to text
         # Exclude the input_ids length to get only the new tokens
         prediction_text = self.tokenizer.decode(prediction_ids[0, input_ids.shape[1]:], skip_special_tokens=True)
+        return [{"generated_text": prediction_text, "ids": prediction_ids[0, input_ids.shape[1]:].tolist()}]
 class StopAtPeriodCriteria(StoppingCriteria):