Fix continue from prefix

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,6 +41,8 @@ class GemmaLlmEngine(LlmEngine):
             stream=True,
             logits_processor=logits_processor_list
         )
+
+        def content_generator():
             for output in stream:
                 choice = output["choices"][0]
                 if choice.get('finish_reason'):
@@ -50,6 +53,8 @@ class GemmaLlmEngine(LlmEngine):
                         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>")