getting generated_ids[0]

Files changed (2) hide show

__pycache__/handler.cpython-311.pyc CHANGED Viewed

Binary files a/__pycache__/handler.cpython-311.pyc and b/__pycache__/handler.cpython-311.pyc differ

handler.py CHANGED Viewed

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

             stopping_criteria=self.stopping_criteria,
         )
+        generated_text = self.tokenizer.decode(generated_ids[0][input_ids.shape[1]:], skip_special_tokens=True)
         prediction = [{"generated_text": generated_text}]
         return prediction