Use Unsloth for QwQ inference

sia/llm_engine/qwq_llm_engine.py

@@ -1,6 +1,9 @@
+# Unsloth should be imported before transformers to ensure all optimizations are applied.
+from unsloth import FastLanguageModel
+
 from pathlib import Path
 from threading import Thread
-from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer, pipeline, BitsAndBytesConfig
+from transformers import AutoTokenizer, TextIteratorStreamer, pipeline
 from typing import Callable, Iterator, Optional
 from xml_schema_validator import XmlLogitsProcessor
 import os
@@ -26,40 +29,27 @@ class QwQLlmEngine(LlmEngine):
             xml_schema_text: Optional XML schema to validate against
         """
         self._temperature = temperature
-                
-        # Configure 4-bit quantization for massive memory savings
-        quantization_config = BitsAndBytesConfig(
-            load_in_4bit=True,
-            bnb_4bit_compute_dtype=torch.bfloat16,
-            bnb_4bit_use_double_quant=True,
-            bnb_4bit_quant_type="nf4"
-        )
 
-        # Load tokenizer first - this uses minimal memory
+        # Load tokenizer
         self._tokenizer = AutoTokenizer.from_pretrained(
             model_path,
-            padding_side="left",
-            trust_remote_code=True,
         )
         
-        # Load model with 4-bit quantization
-        self._model = AutoModelForCausalLM.from_pretrained(
+        # Load model
+        self._model, _returned_tokenizer = FastLanguageModel.from_pretrained(
             model_path,
-            device_map="auto",
-            quantization_config=quantization_config,
-            torch_dtype=torch.bfloat16,
-            low_cpu_mem_usage=True,
-            trust_remote_code=True,
+            load_in_4bit = True, # False for LoRA 16bit
+            fast_inference = True, # Enable vLLM fast inference
+            gpu_memory_utilization = 0.8, # Reduce if out of memory
+            tokenizer = self._tokenizer,
         )
         
-        # Create inference pipeline with memory-efficient settings
+        # Create inference pipeline
         self._pipeline = pipeline(
             "text-generation",
             model=self._model,
             tokenizer=self._tokenizer,
             return_full_text=False,
-            device_map="auto",
-            torch_dtype=torch.bfloat16,
         )
 
         if xml_schema_text:
@@ -102,7 +92,6 @@ class QwQLlmEngine(LlmEngine):
             "text_inputs": text,
             "do_sample": True,
             "temperature": self._temperature,
-            "max_new_tokens": self.token_limit(),
             "streamer": streamer,
             "use_cache": True,
         }