Replaced deepseek with qwq

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:
-            self._llms['deepseek'] = DeepSeekLlmEngine(
-                config.deepseek_model,
-                config.deepseek_temperature,
-                config.deepseek_token_limit,
+        if config.qwq_enabled:
+            self._llms['qwq'] = QwQLlmEngine(
+                config.qwq_model,
+                config.qwq_temperature,
+                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),
-            help='Enable DeepSeek LLM engine (env: SIA_DEEPSEEK_ENABLED)'
+            default=self._parse_bool_env('SIA_QWQ_ENABLED', True),  # Enable by default
+            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'),
-            help='Path to fine-tuned DeepSeek model (env: SIA_DEEPSEEK_MODEL)'
+            default=os.getenv('SIA_QWQ_MODEL', '/root/models/current'),
+            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')),
-            help='DeepSeek temperature (default: 0.6, env: SIA_DEEPSEEK_TEMPERATURE)'
+            default=float(os.getenv('SIA_QWQ_TEMPERATURE', '0.6')),
+            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')),
-            help='DeepSeek token limit (0 for model default, env: SIA_DEEPSEEK_TOKEN_LIMIT)'
+            default=int(os.getenv('SIA_QWQ_TOKEN_LIMIT', '0')),
+            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:
-        return self.args.deepseek_enable
+    def qwq_enabled(self) -> bool:
+        return self.args.qwq_enable
 
     @property
-    def deepseek_model(self) -> str:
-        return self.args.deepseek_model
+    def qwq_model(self) -> str:
+        return self.args.qwq_model
 
     @property
-    def deepseek_temperature(self) -> float:
-        return self.args.deepseek_temperature
+    def qwq_temperature(self) -> float:
+        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