Replaced deepseek with qwq

.gitignore

@@ -3,3 +3,4 @@ __pycache__/
 data/
 model/
 **.egg-info/
+collect.txt

procedures/self_improvement/reasoning.md

@@ -309,25 +309,21 @@ This preserves the temporal relationships between entries while anchoring them t
 
 ## Training Configuration
 
-SIA takes a modular approach to model training by having separate specialized tools for each provider like train_mistral, train_deepseek, etc.
-Each tool shares similar core functionality while handling provider-specific requirements.
-The default training tool and parameters are called from the `/root/sia/tools/train/train.sh` script.
-
 While the training process is conceptually similar across providers, each has unique requirements for data formatting, API interactions, and job management.
-By creating dedicated tools, we can properly encapsulate these differences without complicating the core training logic.
+A dedicated `train` tool encapsulates these differences without complicating the surrounding training logic.
-For example, Mistral needs JSONL files with specific message structures, while other providers might require different formats or metadata.
 
-Training configuration should be consistent regardless of the provider.
+Training configuration is consistent regardless of the provider.
-All training tools read from the same config.yaml format, which defines essential parameters like the system prompt, action schema, and training data paths.
+The same config.yaml file is supported by all implementations.
+This defines essential parameters like the system prompt, action schema, and training data paths.
 These parameters represent fundamental aspects of how we want the model to behave, independent of which provider handles the actual training.
-The tools then translate these standard parameters into provider-specific settings.
+The tool then translates these standard parameters into provider-specific settings for the current active provider.
 
-Training tools enforce important safeguards around version control.
+The training tool enforces important safeguards around version control.
-Before starting a training run, each tool verifies that all source files - including the config itself, training data, system prompt, and action schema - are committed to git.
+Before starting a training run, the tool verifies that all source files are committed to git.
 This ensures reproducibility by guaranteeing we can recreate the exact training conditions that produced any given model.
 The git commit hash becomes part of the internal tracking of model versions.
 
-The tools follow a common workflow:
+The tool follow a common workflow for each provider:
 1. Read and validate the standard config.yaml format
 2. Check that all source files are committed to git
 3. Convert training data into the provider's required format

sia/__main__.py

@@ -3,12 +3,12 @@ import asyncio
 
 from .auto_approver import AutoApprover
 from .config import Config
-from .llm_engine.hf_llm_engine import HfLlmEngine
-from .llm_engine.deepseek_llm_engine import DeepSeekLlmEngine
 from .iteration_logger import IterationLogger
+from .llm_engine.hf_llm_engine import HfLlmEngine
 from .llm_engine.local_llm_engine import LocalLlmEngine
 from .llm_engine.mistral_llm_engine import MistralLlmEngine
 from .llm_engine.openai_llm_engine import OpenAILlmEngine
+from .llm_engine.qwq_llm_engine import QwQLlmEngine
 from .response_parser import ResponseParser
 from .system_metrics import SystemMetrics
 from .web.api import Api
@@ -62,11 +62,11 @@ class Main:
                 config.mistral_api_key,
             )
 
-        if config.deepseek_enabled:
+        if config.qwq_enabled:
-            self._llms['deepseek'] = DeepSeekLlmEngine(
+            self._llms['qwq'] = QwQLlmEngine(
-                config.deepseek_model,
+                config.qwq_model,
-                config.deepseek_temperature,
+                config.qwq_temperature,
-                config.deepseek_token_limit,
+                config.qwq_token_limit,
                 config.hf_api_key,
             )
 

sia/config.py

@@ -184,29 +184,30 @@ class Config:
             default=os.getenv('SIA_MISTRAL_API_KEY'),
             help='Mistral API key (env: SIA_MISTRAL_API_KEY)'
         )
+        # QwQ configuration
         parser.add_argument(
-            '--deepseek-enable',
+            '--qwq-enable',
             action='store_true',
-            default=self._parse_bool_env('SIA_DEEPSEEK_ENABLED', False),
+            default=self._parse_bool_env('SIA_QWQ_ENABLED', True),  # Enable by default
-            help='Enable DeepSeek LLM engine (env: SIA_DEEPSEEK_ENABLED)'
+            help='Enable QwQ LLM engine (env: SIA_QWQ_ENABLED)'
         )
         parser.add_argument(
-            '--deepseek-model',
+            '--qwq-model',
             type=str,
-            default=os.getenv('SIA_DEEPSEEK_MODEL', '/root/models/current'),
+            default=os.getenv('SIA_QWQ_MODEL', '/root/models/current'),
-            help='Path to fine-tuned DeepSeek model (env: SIA_DEEPSEEK_MODEL)'
+            help='Path to QwQ model or HF model ID (default: /root/models/current, env: SIA_QWQ_MODEL)'
         )
         parser.add_argument(
-            '--deepseek-temperature',
+            '--qwq-temperature',
             type=float,
-            default=float(os.getenv('SIA_DEEPSEEK_TEMPERATURE', '0.6')),
+            default=float(os.getenv('SIA_QWQ_TEMPERATURE', '0.6')),
-            help='DeepSeek temperature (default: 0.6, env: SIA_DEEPSEEK_TEMPERATURE)'
+            help='QwQ temperature (default: 0.6, env: SIA_QWQ_TEMPERATURE)'
         )
         parser.add_argument(
-            '--deepseek-token-limit',
+            '--qwq-token-limit',
             type=int,
-            default=int(os.getenv('SIA_DEEPSEEK_TOKEN_LIMIT', '0')),
+            default=int(os.getenv('SIA_QWQ_TOKEN_LIMIT', '0')),
-            help='DeepSeek token limit (0 for model default, env: SIA_DEEPSEEK_TOKEN_LIMIT)'
+            help='QwQ token limit (0 for model default, env: SIA_QWQ_TOKEN_LIMIT)'
         )
 
         self.args = parser.parse_args()
@@ -337,19 +338,20 @@ class Config:
     def mistral_api_key(self) -> Optional[str]:
         return self.args.mistral_api_key
 
+    # QwQ properties
     @property
-    def deepseek_enabled(self) -> bool:
+    def qwq_enabled(self) -> bool:
-        return self.args.deepseek_enable
+        return self.args.qwq_enable
 
     @property
-    def deepseek_model(self) -> str:
+    def qwq_model(self) -> str:
-        return self.args.deepseek_model
+        return self.args.qwq_model
 
     @property
-    def deepseek_temperature(self) -> float:
+    def qwq_temperature(self) -> float:
-        return self.args.deepseek_temperature
+        return self.args.qwq_temperature
 
     @property
-    def deepseek_token_limit(self) -> Optional[int]:
+    def qwq_token_limit(self) -> Optional[int]:
         # Return None if 0 to use model default
-        return self.args.deepseek_token_limit if self.args.deepseek_token_limit > 0 else None
+        return self.args.qwq_token_limit if self.args.qwq_token_limit > 0 else None

sia/llm_engine/deepseek_llm_engine.py

@@ -1,161 +0,0 @@
-from typing import Callable, Iterator, Optional
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer, BitsAndBytesConfig
-from threading import Thread
-from pathlib import Path
-
-from . import LlmEngine
-from .. import util
-
-class DeepSeekLlmEngine(LlmEngine):
-    """
-    LLM Engine implementation for DeepSeek models.
-    Supports fine-tuned DeepSeek-R1 and its distilled versions.
-    """
-    
-    def __init__(
-        self,
-        model_path: str,
-        temperature: float = 0.6,
-        token_limit: Optional[int] = None,
-        api_key: Optional[str] = None,
-    ):
-        """
-        Initialize the DeepSeek LLM Engine.
-
-        Args:
-            model_path: Local path to the fine-tuned model
-            temperature: Sampling temperature (0.6 default as recommended)
-            token_limit: Maximum tokens to generate or context length override
-            api_key: HuggingFace API token if needed
-        """
-        self._model_path = Path(model_path)
-        self._temperature = temperature
-        self._token_limit = token_limit
-        
-        # Load tokenizer with trust_remote_code for DeepSeek models
-        self._tokenizer = AutoTokenizer.from_pretrained(
-            self._model_path,
-            token=api_key,
-            trust_remote_code=True,
-        )
-        
-        # Set padding token to avoid warnings
-        if self._tokenizer.pad_token is None:
-            self._tokenizer.pad_token = self._tokenizer.eos_token
-        
-        # Configure 4-bit quantization with CPU offloading
-        quantization_config = BitsAndBytesConfig(
-            load_in_4bit=True,
-            bnb_4bit_compute_dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16,
-            bnb_4bit_use_double_quant=True,
-            bnb_4bit_quant_type="nf4",
-            llm_int8_enable_fp32_cpu_offload=True
-        )
-        
-        # Configure device map for efficient memory usage
-        # "auto" with the proper quantization config will handle the memory constraints
-        self._device_map = "auto"
-        
-        # Load model with quantization config
-        self._model = AutoModelForCausalLM.from_pretrained(
-            self._model_path,
-            return_dict=True,
-            low_cpu_mem_usage=True,
-            trust_remote_code=True,
-            device_map=self._device_map,
-            quantization_config=quantization_config,
-            torch_dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16,
-            token=api_key,
-        )
-        
-        # Ensure model is in evaluation mode
-        self._model.eval()
-                        
-    def infer(self, system_prompt: str, main_context: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]:
-        """
-        Run inference using the system prompt and main context.
-
-        Args:
-            system_prompt: The system prompt string
-            main_context: The main context string after templating
-            should_stop: Callback that returns True when inference should stop
-
-        Returns:
-            Iterator[str]: An iterator that yields the generated text.
-        """
-        # Tokenize input
-        inputs = self._tokenizer(system_prompt + "\n\n" + main_context, return_tensors="pt").to(self._device_map)
-        
-        # Create streamer for token-by-token generation
-        streamer = TextIteratorStreamer(
-            self._tokenizer,
-            skip_prompt=True,
-            timeout=15.0
-        )
-        
-        # Generate in a separate thread to enable streaming
-        generation_kwargs = {
-            "input_ids": inputs.input_ids,
-            "attention_mask": inputs.attention_mask,
-            "max_new_tokens": self.token_limit() if self._token_limit else 2048,
-            "temperature": self._temperature,
-            "do_sample": True,
-            "streamer": streamer,
-            "repetition_penalty": 1.1,
-            "pad_token_id": self._tokenizer.pad_token_id,
-        }
-        
-        generation_thread = Thread(target=self._model.generate, kwargs=generation_kwargs)
-        generation_thread.start()
-        
-        # Yield tokens as they become available
-        try:
-            for text in streamer:
-                yield text
-                if should_stop():
-                    break
-        finally:
-            # Ensure thread is properly joined even if iteration is interrupted
-            generation_thread.join()
-
-    def token_count(self, system_prompt: str, main_context: str) -> int:
-        """
-        Count tokens for the given system prompt and main context.
-        
-        Args:
-            system_prompt: The system prompt string
-            main_context: The main context string
-            
-        Returns:
-            int: Total number of tokens
-        """
-        combined_prompt = f"{system_prompt}\n\n{main_context}"
-        return len(self._tokenizer.encode(combined_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:
-            config_file = self._model_path / "config.json"
-            if config_file.exists():
-                import json
-                with open(config_file, 'r') as f:
-                    config = json.load(f)
-                if 'max_position_embeddings' in config:
-                    return config['max_position_embeddings']
-                if 'model_max_length' in config:
-                    return config['model_max_length']
-        except Exception:
-            pass
-            
-        # Default to 8k if we can't determine
-        return 8192

sia/llm_engine/qwq_llm_engine.py

@@ -0,0 +1,303 @@
+from typing import Callable, Iterator, Optional
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
+from threading import Thread
+from pathlib import Path
+import sys
+import gc
+import os
+import re
+
+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,
+        temperature: float = 0.6,  # QwQ recommended default
+        token_limit: Optional[int] = None,
+        api_key: Optional[str] = None,
+    ):
+        """
+        Initialize the QwQ LLM Engine.
+
+        Args:
+            model_path: Local path to the model or HF model ID
+            temperature: Sampling temperature (0.6 default as recommended for QwQ)
+            token_limit: Maximum tokens to generate or context length override
+            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 
+        self._top_p = 0.95  # QwQ recommended
+        self._min_p = 0.0   # QwQ recommended
+        self._top_k = 40    # QwQ recommended
+
+        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
+            if self._tokenizer.pad_token is None:
+                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,
+            )
+            
+            # Ensure model is in evaluation mode
+            self._model.eval()
+            print("QwQ model loaded successfully.")
+            
+            # Clear cache after loading
+            gc.collect()
+            
+        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]:
+        """
+        Run inference using the system prompt and main context.
+
+        Args:
+            system_prompt: The system prompt string
+            main_context: The main context string after templating
+            should_stop: Callback that returns True when inference should stop
+
+        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_kwargs = {
+                "input_ids": inputs.input_ids,
+                "attention_mask": inputs.attention_mask,
+                "max_new_tokens": self.token_limit(),
+                "temperature": self._temperature,
+                "top_p": self._top_p,
+                "top_k": self._top_k,
+                "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
+            raw_output = ""
+            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:
+        """
+        Count tokens for the given system prompt and main context.
+        
+        Args:
+            system_prompt: The system prompt string
+            main_context: The main context string
+            
+        Returns:
+            int: Total number of tokens
+        """
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": main_context}
+        ]
+        text = self._tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True,
+        )
+        return len(self._tokenizer.encode(text))
+
+    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/bin/train

@@ -12,4 +12,4 @@ fi
 
 mkdir -p "$OUTPUT_DIR"
 
-train_deepseek --output-dir "$OUTPUT_DIR" --device cpu
+python -m train.qwq "$@"

tools/train/bin/train_deepseek

@@ -1,10 +0,0 @@
-#!/root/venvs/train/bin/python
-"""
-Command-line utility for fine-tuning DeepSeek models using Unsloth.
-Always trains from a base model to create a new fine-tuned model.
-"""
-import sys
-from train.unsloth_deepseek import main
-
-if __name__ == "__main__":
-    sys.exit(main())

tools/train/bin/train_mistral

@@ -1,9 +0,0 @@
-#!/root/venvs/train/bin/python
-"""
-Command-line utility for fine-tuning Mistral models using Mistral API.
-"""
-import sys
-from train.mistral_api import main
-
-if __name__ == "__main__":
-    sys.exit(main())

tools/train/setup.py

@@ -5,26 +5,27 @@ setup(
     version="0.1.0",
     packages=find_packages(),
     scripts=[
-        'bin/train_deepseek',
+        'bin/train'
-        'bin/train_mistral'
     ],
+
     install_requires=[
-        'pyyaml>=6.0',
-        'requests>=2.28.0',
-        'torch>=2.0.0',
-        'transformers>=4.30.0',
         'accelerate>=0.25.0',
         'bitsandbytes>=0.41.1',
-        'einops>=0.7.0',
-        'sentencepiece>=0.1.99',
-        'unsloth>=2025.2',
-        'trl>=0.7.8',
-        'datasets>=2.14.6',
-        'peft>=0.8.0',
-        'pytest>=7.0.0',
-        'pytest-cov>=4.0.0',
         'black>=22.0.0',
-        'flake8>=4.0.0'
+        'datasets>=2.14.6',
+        'einops>=0.7.0',
+        'flake8>=4.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',
+        'transformers>=4.30.0',
+        'trl>=0.7.8',
+        'unsloth>=2025.2',
     ],
     classifiers=[
         'Development Status :: 3 - Alpha',

tools/train/train/qwq.py

@@ -0,0 +1,334 @@
+#!/root/venvs/train/bin/python
+"""
+Fine-tuning script for QwQ models to support SIA's action schema.
+Supports both full and LoRA finetuning methods.
+"""
+import argparse
+import os
+import sys
+import torch
+from dataclasses import dataclass
+from pathlib import Path
+import json
+import logging
+import gc
+
+# Set up logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+# Import from shared library
+from .util import prepare_training_data
+
+@dataclass
+class Config:
+    def __init__(self):
+        parser = argparse.ArgumentParser(description='Train SIA model using QwQ')
+        parser.add_argument(
+            '--config',
+            type=Path,
+            default=Path('/root/sia/training/config.yaml'),
+            help='Path to config file'
+        )
+        parser.add_argument(
+            '--base-model',
+            type=str,
+            default='Qwen/QwQ-32B',
+            help='HuggingFace model ID for base model'
+        )
+        parser.add_argument(
+            '--output-dir',
+            type=Path,
+            required=True,
+            help='Directory to save the trained model'
+        )
+        parser.add_argument(
+            '--api-key',
+            type=str,
+            default=os.environ.get('SIA_HF_API_KEY'),
+            help='HuggingFace API key'
+        )
+        parser.add_argument(
+            '--method',
+            type=str,
+            choices=['lora', 'qlora', 'full'],
+            default='qlora',
+            help='Finetuning method: LoRA, QLoRA (quantized LoRA), or full-model'
+        )
+        parser.add_argument(
+            '--device',
+            type=str,
+            default='auto',
+            help='Override device (cpu, cuda, auto) from config'
+        )
+        self.args = parser.parse_args()
+
+    @property
+    def config_path(self) -> Path:
+        return self.args.config
+
+    @property
+    def base_model(self) -> str:
+        return self.args.base_model
+    
+    @property
+    def output_dir(self) -> Path:
+        return self.args.output_dir
+    
+    @property
+    def api_key(self) -> str:
+        return self.args.api_key
+        
+    @property
+    def device(self) -> str:
+        return self.args.device
+        
+    @property
+    def method(self) -> str:
+        return self.args.method
+
+def format_data_for_qwq(training_data):
+    """
+    Format training data for QwQ model focusing on action schema formats.
+    Ensures each example shows the model how to directly use action elements.
+    """
+    formatted_data = []
+    
+    for sample in training_data:
+        # Get the system prompt, context, and response
+        system_content = ""
+        context_content = ""
+        response_content = ""
+        
+        for message in sample.get("messages", []):
+            if message["role"] == "system":
+                system_content = message["content"]
+            elif message["role"] == "user":
+                context_content = message["content"]
+            elif message["role"] == "assistant":
+                response_content = message["content"]
+        
+        # Create conversations with explicit instruction to use action schema
+        formatted_data.append({
+            "conversations": [
+                {"role": "system", "content": system_content},
+                {"role": "user", "content": context_content},
+                {"role": "assistant", "content": response_content}
+            ]
+        })
+    
+    logger.info(f"Formatted {len(formatted_data)} examples for QwQ training")
+    return formatted_data
+
+def train_model_lora(config, training_data, train_params):
+    """
+    Train QwQ model using LoRA or QLoRA for parameter-efficient fine-tuning.
+    This is the recommended approach for most use cases.
+    """
+    try:
+        # Import required libraries
+        from transformers import (
+            AutoModelForCausalLM, AutoTokenizer, 
+            TrainingArguments, DataCollatorForSeq2Seq
+        )
+        from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
+        from datasets import Dataset
+        from trl import SFTTrainer
+
+    except ImportError as e:
+        logger.error(f"Error importing required libraries: {e}")
+        logger.error("Please ensure transformers, peft, and trl are installed.")
+        sys.exit(1)
+
+    # Format data specifically for QwQ
+    formatted_data = format_data_for_qwq(training_data)
+    dataset = Dataset.from_list(formatted_data)
+
+    logger.info(f"Starting QwQ fine-tuning using {config.method}")
+    logger.info(f"Base model: {config.base_model}")
+    logger.info(f"Device: {config.device}")
+
+    # Configure device mapping and precision
+    if torch.cuda.is_available() and config.device != "cpu":
+        logger.info("Using GPU for training")
+        device_map = "auto"
+        
+        # Configure precision based on method
+        if config.method == "qlora":
+            dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
+            load_in_4bit = True
+            load_in_8bit = False
+        else:
+            dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
+            load_in_4bit = False
+            load_in_8bit = False
+    else:
+        logger.info("Using CPU for training")
+        device_map = "cpu"
+        dtype = torch.float32
+        load_in_4bit = False
+        load_in_8bit = False
+
+    # Configure quantization for QLoRA
+    if config.method == "qlora":
+        from transformers import BitsAndBytesConfig
+        logger.info("Setting up 4-bit quantization for QLoRA")
+        compute_dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
+        
+        bnb_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_compute_dtype=compute_dtype,
+            bnb_4bit_use_double_quant=True
+        )
+    else:
+        bnb_config = None
+
+    # Load tokenizer
+    logger.info(f"Loading tokenizer from {config.base_model}")
+    tokenizer = AutoTokenizer.from_pretrained(
+        config.base_model, 
+        token=config.api_key,
+        trust_remote_code=True
+    )
+    
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    
+    # Load model
+    logger.info(f"Loading model from {config.base_model}")
+    model = AutoModelForCausalLM.from_pretrained(
+        config.base_model,
+        torch_dtype=dtype,
+        device_map=device_map,
+        quantization_config=bnb_config,
+        token=config.api_key,
+        trust_remote_code=True
+    )
+    
+    # Configure LoRA
+    if config.method in ["lora", "qlora"]:
+        if config.method == "qlora":
+            model = prepare_model_for_kbit_training(model)
+            
+        logger.info("Setting up LoRA configuration")
+        target_modules = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
+        
+        lora_config = LoraConfig(
+            r=16,
+            lora_alpha=32,
+            target_modules=target_modules,
+            lora_dropout=0.05,
+            bias="none",
+            task_type="CAUSAL_LM"
+        )
+        
+        model = get_peft_model(model, lora_config)
+        
+    # Create output directory
+    config.output_dir.mkdir(parents=True, exist_ok=True)
+    
+    # Configure training arguments
+    batch_size = train_params.per_device_batch_size
+    gradient_accumulation = train_params.gradient_accumulation_steps
+    
+    # Scale down batch size based on model
+    if "32B" in config.base_model and batch_size > 1:
+        batch_size = 1
+        gradient_accumulation *= 2
+    
+    training_args = TrainingArguments(
+        output_dir=str(config.output_dir),
+        per_device_train_batch_size=batch_size,
+        gradient_accumulation_steps=gradient_accumulation,
+        learning_rate=train_params.learning_rate,
+        num_train_epochs=train_params.epochs,
+        logging_steps=10,
+        save_strategy="epoch",
+        save_total_limit=2,
+        fp16=dtype == torch.float16,
+        bf16=dtype == torch.bfloat16,
+        report_to="none",
+        remove_unused_columns=False,
+        optim="adamw_torch",
+        weight_decay=0.01,
+        max_grad_norm=0.3,
+        warmup_ratio=0.03,
+        lr_scheduler_type="cosine",
+        seed=42
+    )
+    
+    # Set up trainer
+    logger.info("Setting up trainer")
+    trainer = SFTTrainer(
+        model=model,
+        args=training_args,
+        train_dataset=dataset,
+        tokenizer=tokenizer,
+        max_seq_length=train_params.max_seq_length,
+        dataset_text_field="conversations",
+        packing=False
+    )
+    
+    # Start training
+    logger.info("Starting training")
+    trainer.train()
+    
+    # Save the final model
+    logger.info(f"Saving model to {config.output_dir}")
+    trainer.save_model(config.output_dir)
+    tokenizer.save_pretrained(config.output_dir)
+    
+    # Create metadata file
+    with open(config.output_dir / "training_info.json", "w") as f:
+        json.dump({
+            "base_model": config.base_model,
+            "method": config.method,
+            "learning_rate": train_params.learning_rate,
+            "epochs": train_params.epochs,
+            "dataset_size": len(dataset),
+            "batch_size": batch_size,
+            "gradient_accumulation": gradient_accumulation
+        }, f, indent=2)
+    
+    logger.info("Training complete!")
+    return True
+
+def main():
+    # Initialize configuration
+    config = Config()
+    
+    # Prepare training data from config
+    training_data, train_params = prepare_training_data(config.config_path)
+    
+    if not training_data:
+        logger.error("No valid training data found. Exiting.")
+        return 1
+    
+    # Force garbage collection
+    gc.collect()
+    
+    # Train using appropriate method
+    if config.method in ["lora", "qlora"]:
+        success = train_model_lora(config, training_data, train_params)
+    else:
+        logger.error(f"Training method '{config.method}' not yet implemented")
+        return 1
+    
+    if success:
+        # Create symlink to current
+        current_link = Path("/root/models/current")
+        if os.path.exists(current_link) or os.path.islink(current_link):
+            os.unlink(current_link)
+        os.symlink(config.output_dir, current_link, target_is_directory=True)
+        
+        logger.info(f"Training complete. Model saved to {config.output_dir}")
+        logger.info(f"Symlink created at {current_link}")
+        
+        return 0
+    else:
+        logger.error("Training failed")
+        return 1
+
+if __name__ == "__main__":
+    sys.exit(main())

tools/train/train/unsloth_deepseek.py

@@ -1,289 +0,0 @@
-#!/root/venvs/train/bin/python
-"""
-Script for fine-tuning DeepSeek models for SIA using Unsloth.
-Training always starts from a base model and creates a new fine-tuned model.
-"""
-import argparse
-import os
-import sys
-import torch
-from dataclasses import dataclass
-from pathlib import Path
-import json
-
-# Import from shared library
-from .util import prepare_training_data
-
-@dataclass
-class Config:
-    def __init__(self):
-        parser = argparse.ArgumentParser(description='Train SIA model using Unsloth')
-        parser.add_argument(
-            '--config',
-            type=Path,
-            default=Path('/root/sia/training/config.yaml'),
-            help='Path to config file'
-        )
-        parser.add_argument(
-            '--base-model',
-            type=str,
-            default='deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B',
-            help='HuggingFace model ID for base model'
-        )
-        parser.add_argument(
-            '--output-dir',
-            type=Path,
-            required=True,
-            help='Directory to save the trained model'
-        )
-        parser.add_argument(
-            '--api-key',
-            type=str,
-            default=os.environ.get('SIA_HF_API_KEY'),
-            help='HuggingFace API key'
-        )
-        parser.add_argument(
-            '--device',
-            type=str,
-            default='auto',
-            help='Override device (cpu, cuda, auto) from config'
-        )
-        self.args = parser.parse_args()
-
-    @property
-    def config_path(self) -> Path:
-        return self.args.config
-
-    @property
-    def base_model(self) -> str:
-        return self.args.base_model
-    
-    @property
-    def output_dir(self) -> Path:
-        return self.args.output_dir
-    
-    @property
-    def api_key(self) -> str:
-        return self.args.api_key
-        
-    @property
-    def device(self) -> str:
-        return self.args.device
-
-def train_model(config: Config, training_data, train_params):
-    """Train the model using Unsloth"""
-    try:
-        from unsloth import FastLanguageModel
-        from transformers import TrainingArguments, DataCollatorForSeq2Seq
-        from trl import SFTTrainer
-        from datasets import Dataset
-        from unsloth.chat_templates import get_chat_template, train_on_responses_only
-    except ImportError as e:
-        print(f"Error importing required libraries: {e}")
-        print("Please ensure Unsloth and its dependencies are installed.")
-        sys.exit(1)
-
-    print(f"Starting training from base model: {config.base_model}")
-    print(f"Using device: {config.device}")
-    print(f"Training configuration:")
-    print(f"  Max sequence length: {train_params.max_seq_length}")
-    print(f"  Quantization: {train_params.quantization}")
-    print(f"  Batch size: {train_params.per_device_batch_size}")
-    print(f"  Gradient accumulation: {train_params.gradient_accumulation_steps}")
-    print(f"  Mixed precision: {train_params.mixed_precision}")
-    
-    # Convert to datasets format
-    dataset = Dataset.from_list(training_data)
-    
-    # Configure device and dtype
-    if train_params.mixed_precision == "bf16" and torch.cuda.is_bf16_supported():
-        dtype = torch.bfloat16
-    else:
-        dtype = torch.float16
-        
-    # Configure quantization settings
-    load_in_4bit = train_params.quantization == "4bit"
-    load_in_8bit = train_params.quantization == "8bit"
-    
-    # Configure device mapping
-    device_map = config.device
-    if config.device == "cpu":
-        # Force CPU even for quantized model
-        bnb_config = None
-        # When on CPU, we should disable quantization
-        load_in_4bit = False
-        load_in_8bit = False
-        dtype = torch.float32
-        print("CPU-only mode: Disabling quantization and using float32")
-    else:
-        # Setup quantization config for GPU
-        from transformers import BitsAndBytesConfig
-        bnb_config = BitsAndBytesConfig(
-            load_in_4bit=load_in_4bit,
-            load_in_8bit=load_in_8bit,
-            bnb_4bit_use_double_quant=True,
-            bnb_4bit_quant_type="nf4",
-            bnb_4bit_compute_dtype=dtype,
-            llm_int8_enable_fp32_cpu_offload=True
-        ) if (load_in_4bit or load_in_8bit) else None
-    
-    # Load the model with appropriate settings
-    try:
-        model, tokenizer = FastLanguageModel.from_pretrained(
-            model_name=config.base_model,
-            max_seq_length=train_params.max_seq_length,
-            dtype=dtype,
-            quantization_config=bnb_config,
-            device_map=device_map,
-            token=config.api_key,
-        )
-    except Exception as e:
-        print(f"Error loading base model: {e}")
-        sys.exit(1)
-    
-    target_modules = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
-    model = FastLanguageModel.get_peft_model(
-        model,
-        r=8 if config.device == "cpu" else 16,  # Lower rank for CPU to save memory
-        target_modules=target_modules,
-        lora_alpha=16,
-        lora_dropout=0,
-        bias="none",
-        # Only use gradient checkpointing for GPU
-        use_gradient_checkpointing="unsloth" if config.device != "cpu" else None,
-        random_state=3407,
-    )
-    
-    # Function to format conversations
-    def formatting_prompts_func(examples):
-        convos = examples["conversations"]
-        texts = [tokenizer.apply_chat_template(convo, tokenize=False, add_generation_prompt=False) for convo in convos]
-        return {"text": texts}
-    
-    # Standardize dataset and format
-    from unsloth.chat_templates import standardize_sharegpt
-    
-    # Add conversations field if not present
-    if "conversations" not in dataset.column_names:
-        if "messages" in dataset.column_names:
-            dataset = dataset.rename_column("messages", "conversations")
-        else:
-            dataset = dataset.map(lambda x: {"conversations": [{"role": "system", "content": x.get("system_prompt", "")}, 
-                                                            {"role": "user", "content": x.get("prompt", "")}, 
-                                                            {"role": "assistant", "content": x.get("response", "")}]})
-    
-    # Standardize format
-    dataset = standardize_sharegpt(dataset)
-    
-    # Apply formatting
-    dataset = dataset.map(formatting_prompts_func, batched=True)
-    
-    # Configure the trainer
-    config.output_dir.mkdir(parents=True, exist_ok=True)
-    
-    # Determine steps or epochs based on dataset size
-    max_steps = -1
-    num_train_epochs = train_params.epochs
-    if len(dataset) < 100:  # Small dataset
-        # Aim for at least 500 steps for small datasets
-        max_steps = 500
-        num_train_epochs = -1
-    
-    # Configure mixed precision settings
-    fp16 = train_params.mixed_precision == "fp16"
-    bf16 = train_params.mixed_precision == "bf16" and torch.cuda.is_bf16_supported()
-    
-    trainer = SFTTrainer(
-        model=model,
-        tokenizer=tokenizer,
-        train_dataset=dataset,
-        dataset_text_field="text",
-        max_seq_length=train_params.max_seq_length,
-        data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer),
-        dataset_num_proc=1 if config.device == "cpu" else 2,
-        packing=False,
-        args=TrainingArguments(
-            per_device_train_batch_size=train_params.per_device_batch_size,
-            gradient_accumulation_steps=train_params.gradient_accumulation_steps,
-            warmup_steps=5,
-            max_steps=max_steps,
-            num_train_epochs=num_train_epochs,
-            learning_rate=train_params.learning_rate,
-            fp16=fp16,
-            bf16=bf16,
-            logging_steps=10,
-            optim="adamw_torch" if config.device == "cpu" else "adamw_8bit",
-            weight_decay=0.01,
-            lr_scheduler_type="linear",
-            seed=3407,
-            output_dir=str(config.output_dir),
-            report_to="none",
-            dataloader_num_workers=0 if config.device == "cpu" else 2,
-            gradient_checkpointing=config.device != "cpu",
-            max_grad_norm=0.3,
-        ),
-    )
-    
-    # Train only on responses
-    trainer = train_on_responses_only(
-        trainer,
-        instruction_part="<|im_start|>user",
-        response_part="<|im_start|>assistant",
-    )
-    
-    # Train the model
-    trainer.train()
-    
-    # Enable inference mode for the model
-    if config.device != "cpu":
-        model = FastLanguageModel.for_inference(model)
-    
-    # Save the model
-    model.save_pretrained(config.output_dir)
-    tokenizer.save_pretrained(config.output_dir)
-    
-    # Create a metadata file with training information
-    with open(config.output_dir / "training_info.json", "w") as f:
-        json.dump({
-            "base_model": config.base_model,
-            "learning_rate": train_params.learning_rate,
-            "epochs": train_params.epochs,
-            "dataset_size": len(dataset),
-            "device": config.device,
-            "training_method": "unsloth",
-            "max_seq_length": train_params.max_seq_length,
-            "quantization": train_params.quantization,
-        }, f, indent=2)
-    
-def main():
-    config = Config()
-    
-    # Prepare training data
-    training_data, train_params = prepare_training_data(config.config_path)
-    
-    if not training_data:
-        print("No valid training data found. Exiting.")
-        return 1
-    
-    # Train the model
-    try:
-        train_model(config, training_data, train_params)
-        
-        # Create symlink to current
-        current_link = Path("/root/models/current")
-        if os.path.exists(current_link) or os.path.islink(current_link):
-            os.unlink(current_link)
-        os.symlink(config.output_dir, current_link, target_is_directory=True)
-        
-        print(f"Training complete. Model saved to {config.output_dir}")
-        print(f"Symlink created at {current_link}")
-        
-        return 0
-    except Exception as e:
-        print(f"Error during training: {e}")
-        import traceback
-        traceback.print_exc()
-        return 1
-
-if __name__ == "__main__":
-    exit(main())

training/config.yaml

@@ -4,11 +4,6 @@ model:
 params:
   learning_rate: 1e-5
   epochs: 3
-  max_seq_length: 1024
-  quantization: "4bit"  # Options: "none", "4bit", "8bit"
-  per_device_batch_size: 1
-  gradient_accumulation_steps: 8
-  mixed_precision: "no"  # Options: "no", "fp16", "bf16"
 data:
   - "/root/sia/training/clean_start/"
   - "/root/sia/training/delete_indicated_entries/"