Replaced deepseek with qwq
This commit is contained in:
Niels Geens
303
sia/llm_engine/qwq_llm_engine.py
Normal file
303
sia/llm_engine/qwq_llm_engine.py
Normal file
@@ -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
|
||||
Reference in New Issue
Block a user