303 lines
11 KiB
Python
303 lines
11 KiB
Python
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 |