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))