Fixed QwQ tokenizer

sia/llm_engine/qwq_llm_engine.py

@@ -2,6 +2,7 @@ from pathlib import Path
 from threading import Thread
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, TextIteratorStreamer, BitsAndBytesConfig
 from typing import Callable, Iterator
+import json
 import torch
 
 from . import LlmEngine
@@ -26,6 +27,9 @@ class QwQLlmEngine(LlmEngine):
         self._temperature = temperature
         self._token_limit = token_limit
 
+        with open('/root/sia/qwq_tokenizer_config.json', 'r') as f:
+            tokenizer_config = json.load(f)
+
         quantization_config = BitsAndBytesConfig(
             load_in_4bit=True,
             bnb_4bit_compute_dtype=torch.bfloat16,
@@ -41,7 +45,10 @@ class QwQLlmEngine(LlmEngine):
             quantization_config=quantization_config,
         )
 
-        self._tokenizer = AutoTokenizer.from_pretrained(model_path)
+        self._tokenizer = AutoTokenizer.from_pretrained(
+            model_path,
+            tokenizer_config=tokenizer_config,
+        )
 
         self._pipline = pipeline(
             "text-generation",
@@ -92,7 +99,7 @@ class QwQLlmEngine(LlmEngine):
             
         generation_thread.start()
 
-        for text in util.stop_before_value(streamer, '<|eot_id|>'):
+        for text in util.stop_before_value(streamer, self._tokenizer.eos_token):
             yield text
             if should_stop():
                 break