From ab349993ba52c5849ced486af9ac2f12eec8821d Mon Sep 17 00:00:00 2001
From: Niels Geens <niels.geens@gmail.com>
Date: Thu, 22 May 2025 20:24:29 +0000
Subject: [PATCH] Larger context and WiP fix to prefix handling

---
 tools/gemma_infer/src/gemma_infer/__main__.py         |  2 +-
 tools/gemma_infer/src/gemma_infer/gemma_llm_engine.py | 10 +++-------
 2 files changed, 4 insertions(+), 8 deletions(-)

diff --git a/tools/gemma_infer/src/gemma_infer/__main__.py b/tools/gemma_infer/src/gemma_infer/__main__.py
index 477c2a3..b46d857 100644
--- a/tools/gemma_infer/src/gemma_infer/__main__.py
+++ b/tools/gemma_infer/src/gemma_infer/__main__.py
@@ -25,7 +25,7 @@ def main():
     parser.add_argument(
         '--token-limit',
         type=int,
-        default=os.getenv('SIA_GEMMA_TOKEN_LIMIT', 4096),
+        default=os.getenv('SIA_GEMMA_TOKEN_LIMIT', 10000),
         help='Token limit (default: 4096, env: SIA_GEMMA_TOKEN_LIMIT)'
     )
 
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 34f077f..5e85458 100644
--- a/tools/gemma_infer/src/gemma_infer/gemma_llm_engine.py
+++ b/tools/gemma_infer/src/gemma_infer/gemma_llm_engine.py
@@ -60,17 +60,13 @@ class GemmaLlmEngine(LlmEngine):
                 "prefix": True,
             },
         ] if prefix else [
-            {
-                "role": "system",
-                "content": system,
-            },
             {
                 "role": "user",
-                "content": context,
-            },
+                "content": f"{system}\n\n--- Context ---\n{context}",
+            }
         ]
         return self._tokenizer.apply_chat_template(
             messages,
             tokenize=False,
-            add_generation_prompt=True
+            add_generation_prompt=prefix is None or len(prefix) == 0
         ).removeprefix("<bos>")
\ No newline at end of file