From 58213913098db8a791719eb685eb88af1e9f28fd Mon Sep 17 00:00:00 2001
From: Niels Geens <niels.geens@gmail.com>
Date: Sat, 10 Jan 2026 13:50:57 +0100
Subject: [PATCH] Fix continue from prefix

---
 .../src/gemma_infer/gemma_llm_engine.py       | 25 +++++++++++--------
 1 file changed, 15 insertions(+), 10 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 fd9f6d1..e688dec 100644
--- a/tools/gemma_infer/src/gemma_infer/gemma_llm_engine.py
+++ b/tools/gemma_infer/src/gemma_infer/gemma_llm_engine.py
@@ -5,6 +5,7 @@ os.chdir("/usr/local/lib")
 
 from llama_cpp import Llama, LogitsProcessorList, llama_cpp
 from llm_engine_utils import LlmEngine
+from llm_engine_utils.iterators import skip_prefix
 from pathlib import Path
 from transformers import AutoTokenizer
 from typing import Iterator
@@ -40,15 +41,19 @@ class GemmaLlmEngine(LlmEngine):
             stream=True,
             logits_processor=logits_processor_list
         )
-        for output in stream:
-            choice = output["choices"][0]
-            if choice.get('finish_reason'):
-                break
-            if 'text' in choice:
-                text = choice['text']
-                if text == '':
+
+        def content_generator():
+            for output in stream:
+                choice = output["choices"][0]
+                if choice.get('finish_reason'):
                     break
-                yield text
+                if 'text' in choice:
+                    text = choice['text']
+                    if text == '':
+                        break
+                    yield text
+
+        yield from skip_prefix(content_generator(), prefix)
 
     def token_count(self, system: str, context: str) -> int:
         return len(self._format_messages(system, context, None))
@@ -65,7 +70,6 @@ class GemmaLlmEngine(LlmEngine):
             {
                 "role": "assistant",
                 "content": prefix,
-                "prefix": True,
             },
         ] if prefix else [
             {
@@ -76,5 +80,6 @@ class GemmaLlmEngine(LlmEngine):
         return self._tokenizer.apply_chat_template(
             messages,
             tokenize=False,
-            add_generation_prompt=prefix is None or len(prefix) == 0
+            add_generation_prompt=not prefix,
+            continue_final_message=bool(prefix)
         ).removeprefix("<bos>")
\ No newline at end of file