From 5e47200ece8a5bbd7a77d220f694617180a719a7 Mon Sep 17 00:00:00 2001
From: Niels Geens <niels.geens@gmail.com>
Date: Fri, 18 Apr 2025 12:30:47 +0000
Subject: [PATCH] Fixes for QwQ llm engine

---
 sia/llm_engine/qwq_llm_engine.py | 4 ++--
 system_prompt.md                 | 6 +++++-
 tools/train/train/qwq.py         | 2 +-
 3 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/sia/llm_engine/qwq_llm_engine.py b/sia/llm_engine/qwq_llm_engine.py
index d628f6e..fc0aa7d 100644
--- a/sia/llm_engine/qwq_llm_engine.py
+++ b/sia/llm_engine/qwq_llm_engine.py
@@ -47,7 +47,7 @@ class QwQLlmEngine(LlmEngine):
 
         self._tokenizer = AutoTokenizer.from_pretrained(
             model_path,
-            tokenizer_config=tokenizer_config,
+            **tokenizer_config,
         )
 
         self._pipeline = pipeline(
@@ -87,7 +87,7 @@ class QwQLlmEngine(LlmEngine):
             tokenize=False,
             add_generation_prompt=False,
         )
-            
+
         streamer = TextIteratorStreamer(
             self._tokenizer,
             skip_prompt=True,
diff --git a/system_prompt.md b/system_prompt.md
index 154c719..eb31b09 100644
--- a/system_prompt.md
+++ b/system_prompt.md
@@ -88,4 +88,8 @@ Use scripts and background processes to do so.
 
 The user may take some time to respond or may forget to respond.
 Keep detailed notes of your interactions and your expectations regarding time!
-Avoid overflowing the user with many messages.
\ No newline at end of file
+Avoid overflowing the user with many messages.
+
+# Action schema
+
+What follows is the xml schema you should adhere to:
diff --git a/tools/train/train/qwq.py b/tools/train/train/qwq.py
index 8a3d3f8..d42ed05 100644
--- a/tools/train/train/qwq.py
+++ b/tools/train/train/qwq.py
@@ -79,7 +79,7 @@ def main():
 
     tokenizer = AutoTokenizer.from_pretrained(
         args.base_model,
-        tokenizer_config=tokenizer_config,
+        **tokenizer_config,
     )
 
     model, _returned_tokenizer = FastLanguageModel.from_pretrained(