Converted QwQ notebooks to .py files

sia/__main__.py

@@ -67,7 +67,6 @@ class Main:
                 config.qwq_model,
                 config.qwq_temperature,
                 config.qwq_token_limit,
-                config.hf_api_key,
             )
 
         if not self._llms:

sia/llm_engine/local_llm_engine.py

@@ -21,8 +21,8 @@ class LocalLlmEngine(LlmEngine):
         Args:
             model_path: Path to the model weights to be used.
             temperature: Temperature for sampling
-            api_token: Huggingface API key
             token_limit: Maximum number of tokens to generate
+            api_token: Huggingface API key
         """
         self._temperature = temperature
         self._token_limit = token_limit

sia/llm_engine/qwq_llm_engine.py

@@ -1,98 +1,56 @@
-from typing import Callable, Iterator, Optional
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
-from threading import Thread
 from pathlib import Path
-import sys
+from threading import Thread
-import gc
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, TextIteratorStreamer, BitsAndBytesConfig
-import os
+from typing import Callable, Iterator
-import re
+import torch
 
 from . import LlmEngine
 from .. import util
 
 class QwQLlmEngine(LlmEngine):
-    """
-    LLM Engine implementation for QwQ models.
-    
-    QwQ is a reasoning-based model with <think> capabilities. This engine handles:
-    1. Proper initialization with recommended parameters
-    2. Processing outputs to extract reasoning and actions
-    3. Converting QwQ's format to SIA-compatible action schemas
-    """
     
     def __init__(
         self,
-        model_path: str,
+        model_path: Path,
-        temperature: float = 0.6,  # QwQ recommended default
+        temperature: float,
-        token_limit: Optional[int] = None,
+        token_limit: int = None,
-        api_key: Optional[str] = None,
     ):
         """
         Initialize the QwQ LLM Engine.
 
         Args:
-            model_path: Local path to the model or HF model ID
+            model_path: Local path to the model
-            temperature: Sampling temperature (0.6 default as recommended for QwQ)
+            temperature: Sampling temperature
-            token_limit: Maximum tokens to generate or context length override
+            token_limit: Maximum tokens to generate
-            api_key: HuggingFace API token if needed
         """
-        self._model_path = Path(model_path) if os.path.exists(model_path) else model_path
         self._temperature = temperature
         self._token_limit = token_limit
 
-        # QwQ-specific parameters 
+        quantization_config = BitsAndBytesConfig(
-        self._top_p = 0.95  # QwQ recommended
+            load_in_4bit=True,
-        self._min_p = 0.0   # QwQ recommended
+            bnb_4bit_compute_dtype=torch.bfloat16,
-        self._top_k = 40    # QwQ recommended
+            bnb_4bit_quant_type="nf4",
-
+            bnb_4bit_use_double_quant=True
-        try:
-            # Free memory before loading
-            gc.collect()
-            
-            print(f"Loading QwQ tokenizer from {self._model_path}...")
-            self._tokenizer = AutoTokenizer.from_pretrained(
-                self._model_path,
-                token=api_key,
-                trust_remote_code=True,
         )
 
-            # Set padding token to avoid warnings
+        model = AutoModelForCausalLM.from_pretrained(
-            if self._tokenizer.pad_token is None:
+            model_path,
-                self._tokenizer.pad_token = self._tokenizer.eos_token
-            
-            # Device configuration
-            if torch.cuda.is_available():
-                print(f"Loading QwQ model on GPU...")
-                device_map = "auto"
-                dtype = torch.bfloat16
-            else:
-                print(f"Loading QwQ model on CPU...")
-                device_map = "cpu"
-                dtype = torch.float32
-            
-            # Load model with appropriate settings
-            self._model = AutoModelForCausalLM.from_pretrained(
-                self._model_path,
-                device_map=device_map,
-                torch_dtype=dtype,
-                trust_remote_code=True,
             return_dict=True,
-                token=api_key,
+            device_map="auto",
+            attn_implementation="flash_attention_2",
+            use_cache=True,
+            quantization_config=quantization_config,
         )
 
-            # Ensure model is in evaluation mode
+        self._tokenizer = AutoTokenizer.from_pretrained(model_path)
-            self._model.eval()
-            print("QwQ model loaded successfully.")
 
-            # Clear cache after loading
+        self._pipline = pipeline(
-            gc.collect()
+            "text-generation",
+            model=model,
+            tokenizer=self._tokenizer,
+            return_full_text=False,
+        )
 
-        except Exception as e:
-            print(f"Failed to initialize QwQ model: {e}")
-            import traceback
-            traceback.print_exc()
-            raise RuntimeError(f"Failed to initialize QwQ model: {e}")
                         
     def infer(self, system_prompt: str, main_context: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]:
         """
@@ -106,142 +64,41 @@ class QwQLlmEngine(LlmEngine):
         Returns:
             Iterator[str]: An iterator that yields the generated text.
         """
-        try:
-            # Format as messages for chat template
         messages = [
             {"role": "system", "content": system_prompt},
             {"role": "user", "content": main_context}
         ]
 
-            # Apply chat template - DO NOT add <think> token as it will be handled by the model
         text = self._tokenizer.apply_chat_template(
             messages,
             tokenize=False,
             add_generation_prompt=True,
         )
             
-            # Tokenize input
-            print("Tokenizing input...")
-            inputs = self._tokenizer(text, return_tensors="pt").to(self._model.device)
-            
-            # Create streamer for token-by-token generation
-            print("Starting generation...")
         streamer = TextIteratorStreamer(
             self._tokenizer,
             skip_prompt=True,
-                skip_special_tokens=True,
-                timeout=60.0
         )
 
-            # Configure generation with QwQ's recommended parameters
+        generation_thread = Thread(
-            generation_kwargs = {
+            target=self._pipline,
-                "input_ids": inputs.input_ids,
+            kwargs=dict(
-                "attention_mask": inputs.attention_mask,
+                text_inputs=text,
-                "max_new_tokens": self.token_limit(),
+                do_sample=True,
-                "temperature": self._temperature,
+                temperature=self._temperature,
-                "top_p": self._top_p,
+                max_new_tokens=self._token_limit,
-                "top_k": self._top_k,
+                streamer=streamer,
-                "min_p": self._min_p,
+            )
-                "do_sample": True,
+        )
-                "streamer": streamer,
-                "repetition_penalty": 1.1,
-                "pad_token_id": self._tokenizer.pad_token_id,
-                "use_cache": True,
-            }
             
-            print("Starting generation thread...")
-            generation_thread = Thread(target=self._model.generate, kwargs=generation_kwargs)
         generation_thread.start()
 
-            # Accumulate raw output and track think mode
+        for text in util.stop_before_value(streamer, '<|eot_id|>'):
-            raw_output = ""
+            yield text
-            action_extracted = False
-            
-            # Process thinking and extract actions
-            try:
-                for text in streamer:
-                    raw_output += text
-                    
-                    # Check if we should stop
             if should_stop():
-                        print("Generation stopped by caller")
                 break
 
-                    # Extract action if available
-                    action = self._extract_action(raw_output)
-                    if action and not action_extracted:
-                        # We've found an action tag - yield it
-                        action_extracted = True
-                        yield action
-                    elif not action_extracted:
-                        # Still in thinking phase or no action yet - yield tokens
-                        yield text
-                
-                # Process remaining output
-                if raw_output and not action_extracted:
-                    final_action = self._process_final_output(raw_output)
-                    if final_action:
-                        yield final_action
-                    
-            finally:
-                # Ensure thread is properly joined even if iteration is interrupted
         generation_thread.join()
-                # Force garbage collection after generation
-                gc.collect()
-                
-        except Exception as e:
-            print(f"QwQ inference error: {e}")
-            import traceback
-            traceback.print_exc()
-            # Re-raise to make the failure visible
-            raise RuntimeError(f"QwQ inference failed: {e}")
-    
-    def _extract_action(self, text: str) -> Optional[str]:
-        """
-        Extract SIA-compatible action from QwQ output.
-        Returns the action if found, None if still in thinking mode.
-        """
-        # Check if we have a complete think block followed by an action
-        think_pattern = r'<think>(.*?)</think>\s*(<\w+.*?>)'
-        match = re.search(think_pattern, text, re.DOTALL)
-        
-        if match:
-            # Found a think block followed by an action tag
-            action_start = match.group(2)
-            # Return the action part
-            action_idx = text.index(action_start)
-            return text[action_idx:]
-        
-        # Check for direct action (no thinking)
-        action_pattern = r'^(<(?:single|repeat|delete|stop|reasoning|read_stdin|write_stdout).*?>)'
-        match = re.search(action_pattern, text)
-        if match:
-            return text
-            
-        return None
-    
-    def _process_final_output(self, text: str) -> str:
-        """
-        Process final output if no action was extracted.
-        Converts thinking content to reasoning if needed.
-        """
-        # Check if there's thinking content
-        think_pattern = r'<think>(.*?)</think>'
-        match = re.search(think_pattern, text, re.DOTALL)
-        
-        if match:
-            # Extract thinking content
-            thinking = match.group(1).strip()
-            if thinking:
-                # Convert to reasoning
-                return f"<reasoning>\n{thinking}\n</reasoning>"
-        
-        # If the response has no XML tags but isn't empty, make it reasoning
-        if text.strip() and not re.search(r'<\w+.*?>', text):
-            return f"<reasoning>\n{text.strip()}\n</reasoning>"
-            
-        # Return as-is if it already has valid XML tags
-        return text
 
     def token_count(self, system_prompt: str, main_context: str) -> int:
         """
@@ -258,46 +115,10 @@ class QwQLlmEngine(LlmEngine):
             {"role": "system", "content": system_prompt},
             {"role": "user", "content": main_context}
         ]
-        text = self._tokenizer.apply_chat_template(
+        prompt = self._tokenizer.apply_chat_template(
-            messages,
+            messages, tokenize=False, add_generation_prompt=True
-            tokenize=False,
-            add_generation_prompt=True,
         )
-        return len(self._tokenizer.encode(text))
+        return len(self._tokenizer.encode(prompt))
 
     def token_limit(self) -> int:
-        """
-        Get the model's context window size.
-        
-        Returns:
-            int: Maximum number of tokens the model can process
-        """
-        if self._token_limit is not None:
         return self._token_limit
-        
-        # Try to detect model size from config
-        try:
-            if isinstance(self._model_path, Path):
-                config_file = self._model_path / "config.json"
-                if config_file.exists():
-                    import json
-                    with open(config_file, 'r') as f:
-                        config = json.load(f)
-                else:
-                    config = self._model.config.to_dict()
-            else:
-                config = self._model.config.to_dict()
-                
-            # Check for context length in different possible fields
-            if 'max_position_embeddings' in config:
-                return config['max_position_embeddings']
-            if 'model_max_length' in config:
-                return config['model_max_length']
-            
-            # Safe fallback for QwQ - it supports up to 8192 by default
-            return 8192
-        except Exception as e:
-            print(f"Warning: Failed to read model config: {e}")
-            
-        # Default fallback
-        return 4096

tools/train/train/qwq.py

@@ -6,7 +6,6 @@ Fine-tuning for QwQ model
 # Unsloth should be imported before transformers to ensure all optimizations are applied.
 from unsloth import FastLanguageModel, is_bfloat16_supported
 
-from .dataset import Dataset
 from dataclasses import dataclass
 from pathlib import Path
 from transformers import TrainingArguments
@@ -15,6 +14,8 @@ from typing import Optional, List
 import argparse
 import os
 
+from .dataset import Dataset
+
 @dataclass
 class Args:
     def __init__(self, args: Optional[List[str]]):
@@ -78,7 +79,7 @@ def main():
         load_in_4bit = True, # False for LoRA 16bit
         fast_inference = True, # Enable vLLM fast inference
         max_lora_rank = lora_rank,
-        gpu_memory_utilization = 0.85, # Reduce if out of memory
+        gpu_memory_utilization = 0.5, # Reduce if out of memory
     )
 
     model = FastLanguageModel.get_peft_model(
@@ -97,12 +98,6 @@ def main():
         loftq_config = None, # And LoftQ
     )
 
-    response_template = tokenizer.apply_chat_template(
-        [{"role": "assistant", "content": ""}],
-        tokenize=False,
-        add_generation_prompt=True
-    )
-
     training_args = TrainingArguments(
         output_dir=str(args.output_dir),
         num_train_epochs=3,
@@ -129,10 +124,6 @@ def main():
         train_dataset=dataset.to_transformers_dataset(tokenizer),
         dataset_text_field="messages",
         max_seq_length=max_seq_length,
-        data_collator=DataCollatorForCompletionOnlyLM(
-            response_template=response_template, 
-            tokenizer=tokenizer
-        ),
     )
 
     trainer.train()
@@ -140,7 +131,6 @@ def main():
     model.save_pretrained_merged(
         str(args.output_dir), 
         tokenizer=tokenizer,
-        save_method="merged_16bit"
     )
 
 if __name__ == "__main__":

tools/train/train/qwq_train.ipynb

@@ -22,7 +22,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [