Fixed QwQ tokenizer

tools/train/train/qwq.py

@@ -8,10 +8,11 @@ from unsloth import FastLanguageModel, is_bfloat16_supported
 
 from dataclasses import dataclass
 from pathlib import Path
-from transformers import TrainingArguments
-from trl import SFTTrainer, DataCollatorForCompletionOnlyLM
+from transformers import AutoTokenizer, TrainingArguments
+from trl import SFTTrainer
 from typing import Optional, List
 import argparse
+import json
 import os
 
 from .dataset import Dataset
@@ -73,13 +74,22 @@ def main():
     max_seq_length = 2048 # Can increase for longer reasoning traces
     lora_rank = 64 # Larger rank = smarter, but slower
 
-    model, tokenizer = FastLanguageModel.from_pretrained(
+    with open('/root/sia/qwq_tokenizer_config.json', 'r') as f:
+        tokenizer_config = json.load(f)
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        args.base_model,
+        tokenizer_config=tokenizer_config,
+    )
+
+    model, _returned_tokenizer = FastLanguageModel.from_pretrained(
         model_name = args.base_model,
         max_seq_length = max_seq_length,
         load_in_4bit = True, # False for LoRA 16bit
         fast_inference = True, # Enable vLLM fast inference
         max_lora_rank = lora_rank,
         gpu_memory_utilization = 0.5, # Reduce if out of memory
+        tokenizer = tokenizer,
     )
 
     model = FastLanguageModel.get_peft_model(