From c1cef08941bbfb1767fd6dc486fecb1493a68664 Mon Sep 17 00:00:00 2001
From: Niels Geens <niels.geens@gmail.com>
Date: Sat, 19 Apr 2025 14:12:00 +0000
Subject: [PATCH] Speed up inference

---
 setup.py                         | 44 +++++++++++++++++---------------
 sia/llm_engine/qwq_llm_engine.py | 42 ++++++++++++++++++------------
 2 files changed, 50 insertions(+), 36 deletions(-)

diff --git a/setup.py b/setup.py
index a35bf98..4f1c9de 100644
--- a/setup.py
+++ b/setup.py
@@ -1,30 +1,34 @@
 from setuptools import setup, find_packages
 
 setup(
-    name="sia",
+    name="train",
     version="0.1.0",
     packages=find_packages(),
-    entry_points={
-        'console_scripts': [
-            'sia=sia.__main__:main',
-        ],
-    },
+    scripts=[
+        'bin/train'
+    ],
+
     install_requires=[
+        'accelerate>=0.25.0',
+        'bitsandbytes>=0.45.0',
+        'black>=22.0.0',
+        'datasets>=2.14.6',
+        'einops>=0.7.0',
+        'flake8>=4.0.0',
+        'ipykernel>=6.0.0', 
+        'ipywidgets>=8.0.0',
+        'peft>=0.8.0',
+        'peft>=0.8.0',
+        'pytest-cov>=4.0.0',
+        'pytest>=7.0.0',
+        'pyyaml>=6.0',
+        'requests>=2.28.0',
+        'sentencepiece>=0.1.99',
         'torch>=2.0.0',
-        'accelerate>=0.26.0',
-        'aiohttp>=3.8.0',
-        'bitsandbytes>=0.41.0',
-        'dotenv-python>=0.0.1',
-        'huggingface_hub>=0.16.0',
-        'lxml>=4.9.0',
-        'mistral-common>=1.0.0',
-        'mistralai>=0.0.7',
-        'openai>=1.0.0',
-        'psutil>=5.9.0',
-        'python-dotenv>=1.0.0',
-        'tiktoken>=0.4.0',
         'transformers>=4.30.0',
-        'xml_schema_validator @ file:///root/sia/lib/xml_schema_validator',
+        'trl>=0.7.8',
+        'unsloth>=2025.3',
+        'vllm==0.8.2',
     ],
     python_requires='>=3.10',
-)
\ No newline at end of file
+)
diff --git a/sia/llm_engine/qwq_llm_engine.py b/sia/llm_engine/qwq_llm_engine.py
index e3c8b71..1c0d0bb 100644
--- a/sia/llm_engine/qwq_llm_engine.py
+++ b/sia/llm_engine/qwq_llm_engine.py
@@ -1,9 +1,9 @@
 from pathlib import Path
 from threading import Thread
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, TextIteratorStreamer, BitsAndBytesConfig
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer, pipeline, BitsAndBytesConfig
 from typing import Callable, Iterator, Optional
 from xml_schema_validator import XmlLogitsProcessor
-import json
+import os
 import torch
 
 from . import LlmEngine
@@ -26,31 +26,40 @@ 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_quant_type="nf4",
-            bnb_4bit_use_double_quant=True
-        )
-
-        model = AutoModelForCausalLM.from_pretrained(
-            model_path,
-            return_dict=True,
-            device_map="auto",
-            use_cache=True,
-            quantization_config=quantization_config,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type="nf4"
         )
 
+        # Load tokenizer first - this uses minimal memory
         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(
+            model_path,
+            device_map="auto",
+            quantization_config=quantization_config,
+            torch_dtype=torch.bfloat16,
+            low_cpu_mem_usage=True,
+            trust_remote_code=True,
+        )
+        
+        # Create inference pipeline with memory-efficient settings
         self._pipeline = pipeline(
             "text-generation",
-            model=model,
+            model=self._model,
             tokenizer=self._tokenizer,
             return_full_text=False,
+            device_map="auto",
+            torch_dtype=torch.bfloat16,
         )
 
         if xml_schema_text:
@@ -95,6 +104,7 @@ class QwQLlmEngine(LlmEngine):
             "temperature": self._temperature,
             "max_new_tokens": self.token_limit(),
             "streamer": streamer,
+            "use_cache": True,
         }
         
         if self._logits_processor:
@@ -135,4 +145,4 @@ class QwQLlmEngine(LlmEngine):
         return len(self._tokenizer.encode(prompt))
 
     def token_limit(self) -> int:
-        return self._pipeline.model.config.max_position_embeddings
\ No newline at end of file
+        return self._pipeline.model.config.max_position_embeddings