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