From d2074cfcaa69ab09ef4f6deb6a4c348a62426121 Mon Sep 17 00:00:00 2001 From: Sindhu Somasundaram <56774226+sindhuvahinis@users.noreply.github.com> Date: Mon, 9 Sep 2024 16:58:46 -0700 Subject: [PATCH] [python] check whether last token is generated for json_output_formatter (#2375) --- engines/python/setup/djl_python/output_formatter.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/engines/python/setup/djl_python/output_formatter.py b/engines/python/setup/djl_python/output_formatter.py index 17709388e..3689a65d6 100644 --- a/engines/python/setup/djl_python/output_formatter.py +++ b/engines/python/setup/djl_python/output_formatter.py @@ -110,8 +110,8 @@ def _json_output_formatter(request_output: TextGenerationOutput): request_output.best_sequence_index] # TODO: Fix this so it is not required. Right now, this call is needed to # advance the token iterator, which is needed for rolling batch to work properly - next_token, _, _ = best_sequence.get_next_token() - if not request_output.finished: + next_token, _, is_last_token = best_sequence.get_next_token() + if not is_last_token: return "" details = get_details_dict(request_output, include_tokens=True) if details.get("finish_reason") == "error":