From 39e0c81ef13c8a31226911652289d6828a802313 Mon Sep 17 00:00:00 2001 From: Niels Geens Date: Sat, 10 Jan 2026 13:41:32 +0100 Subject: [PATCH] Better stop token handling --- tools/gemma_infer/src/gemma_infer/gemma_llm_engine.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/tools/gemma_infer/src/gemma_infer/gemma_llm_engine.py b/tools/gemma_infer/src/gemma_infer/gemma_llm_engine.py index d06fd11..fd9f6d1 100644 --- a/tools/gemma_infer/src/gemma_infer/gemma_llm_engine.py +++ b/tools/gemma_infer/src/gemma_infer/gemma_llm_engine.py @@ -41,8 +41,14 @@ class GemmaLlmEngine(LlmEngine): logits_processor=logits_processor_list ) for output in stream: - if 'text' in output["choices"][0]: - yield output["choices"][0]['text'] + choice = output["choices"][0] + if choice.get('finish_reason'): + break + if 'text' in choice: + text = choice['text'] + if text == '': + break + yield text def token_count(self, system: str, context: str) -> int: return len(self._format_messages(system, context, None))