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 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 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'(.*?)\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'(.*?)' 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"\n{thinking}\n" # 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"\n{text.strip()}\n" # 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