llm/SIA
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Converted QwQ notebooks to .py files

Browse Source
This commit is contained in:
Niels Geens
2025-04-01 10:41:24 +02:00
parent 642c5181bb
commit 6f3d414d17
5 changed files with 368 additions and 558 deletions
Download Patch File Download Diff File Expand all files Collapse all files

239
sia/__main__.py
View File

@@ -1,121 +1,120 @@
from aiohttp import web
import asyncio
from .auto_approver import AutoApprover
from .config import Config
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
from .web.static import Static
from .web.websockts import Websockets
from .web_agent import WebAgent
from .web_io_buffer import WebIOBuffer
from .working_memory import WorkingMemory
from .xml_validator import XMLValidator
class Main:
@classmethod
async def create(cls, config: Config):
self = cls()
self._config = config
self._system_prompt = self._config.system_prompt.read_text()
self._action_schema = self._config.action_schema.read_text()
# Initialize LLM engines based on config
self._llms = {}
if config.local_enabled:
self._llms['local'] = LocalLlmEngine(
config.local_model,
config.local_temperature,
config.local_token_limit,
config.local_api_key,
)
if config.openai_enabled:
self._llms['openai'] = OpenAILlmEngine(
config.openai_model,
config.openai_temperature,
config.openai_token_limit,
config.openai_api_key,
)
if config.hf_enabled:
self._llms['hf'] = HfLlmEngine(
config.hf_model,
config.hf_temperature,
config.hf_api_key,
)
if config.mistral_enabled:
self._llms['mistral'] = MistralLlmEngine(
config.mistral_model,
config.mistral_temperature,
config.mistral_token_limit,
config.mistral_api_key,
)
if config.qwq_enabled:
self._llms['qwq'] = QwQLlmEngine(
config.qwq_model,
config.qwq_temperature,
config.qwq_token_limit,
config.hf_api_key,
)
if not self._llms:
raise ValueError("No LLM engines enabled in configuration")
self._io_buffer = WebIOBuffer()
self._working_memory = WorkingMemory()
self._agent = WebAgent(
system_prompt=self._system_prompt,
action_schema=self._action_schema,
working_memory=self._working_memory,
metrics=SystemMetrics(),
llms=self._llms,
validator=XMLValidator(self._action_schema),
parser=ResponseParser(config.work_dir, self._io_buffer),
iteration_logger=IterationLogger(self._config.iterations_dir, self._system_prompt, self._action_schema),
)
self._auto_approver = AutoApprover(self._agent)
self._app = web.Application()
self._api = Api(config.work_dir, self._app, self._agent, self._io_buffer, self._working_memory, self._auto_approver)
self._websockets = Websockets(self._app, self._agent, self._io_buffer, self._auto_approver, self._working_memory)
self._static = Static(self._app, self._config)
return self
@property
def app(self):
return self._app
async def _serve_index(self, request: web.Request) -> web.Response:
"""Serve the React application HTML for any unmatched routes."""
index_path = self._config.static_files / "index.html"
if not index_path.exists():
raise web.HTTPNotFound()
with open(index_path, "r") as f:
html_content = f.read()
return web.Response(
text=html_content,
content_type="text/html"
)
def main():
loop = asyncio.new_event_loop()
config = Config()
main_instance = loop.run_until_complete(Main.create(config))
print(f"Web server started at http://localhost:{config.port}")
web.run_app(main_instance.app, loop=loop, host=config.host, port=config.port)
from aiohttp import web
import asyncio
from .auto_approver import AutoApprover
from .config import Config
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
from .web.static import Static
from .web.websockts import Websockets
from .web_agent import WebAgent
from .web_io_buffer import WebIOBuffer
from .working_memory import WorkingMemory
from .xml_validator import XMLValidator
class Main:
@classmethod
async def create(cls, config: Config):
self = cls()
self._config = config
self._system_prompt = self._config.system_prompt.read_text()
self._action_schema = self._config.action_schema.read_text()
# Initialize LLM engines based on config
self._llms = {}
if config.local_enabled:
self._llms['local'] = LocalLlmEngine(
config.local_model,
config.local_temperature,
config.local_token_limit,
config.local_api_key,
)
if config.openai_enabled:
self._llms['openai'] = OpenAILlmEngine(
config.openai_model,
config.openai_temperature,
config.openai_token_limit,
config.openai_api_key,
)
if config.hf_enabled:
self._llms['hf'] = HfLlmEngine(
config.hf_model,
config.hf_temperature,
config.hf_api_key,
)
if config.mistral_enabled:
self._llms['mistral'] = MistralLlmEngine(
config.mistral_model,
config.mistral_temperature,
config.mistral_token_limit,
config.mistral_api_key,
)
if config.qwq_enabled:
self._llms['qwq'] = QwQLlmEngine(
config.qwq_model,
config.qwq_temperature,
config.qwq_token_limit,
)
if not self._llms:
raise ValueError("No LLM engines enabled in configuration")
self._io_buffer = WebIOBuffer()
self._working_memory = WorkingMemory()
self._agent = WebAgent(
system_prompt=self._system_prompt,
action_schema=self._action_schema,
working_memory=self._working_memory,
metrics=SystemMetrics(),
llms=self._llms,
validator=XMLValidator(self._action_schema),
parser=ResponseParser(config.work_dir, self._io_buffer),
iteration_logger=IterationLogger(self._config.iterations_dir, self._system_prompt, self._action_schema),
)
self._auto_approver = AutoApprover(self._agent)
self._app = web.Application()
self._api = Api(config.work_dir, self._app, self._agent, self._io_buffer, self._working_memory, self._auto_approver)
self._websockets = Websockets(self._app, self._agent, self._io_buffer, self._auto_approver, self._working_memory)
self._static = Static(self._app, self._config)
return self
@property
def app(self):
return self._app
async def _serve_index(self, request: web.Request) -> web.Response:
"""Serve the React application HTML for any unmatched routes."""
index_path = self._config.static_files / "index.html"
if not index_path.exists():
raise web.HTTPNotFound()
with open(index_path, "r") as f:
html_content = f.read()
return web.Response(
text=html_content,
content_type="text/html"
)
def main():
loop = asyncio.new_event_loop()
config = Config()
main_instance = loop.run_until_complete(Main.create(config))
print(f"Web server started at http://localhost:{config.port}")
web.run_app(main_instance.app, loop=loop, host=config.host, port=config.port)
return 0

242
sia/llm_engine/local_llm_engine.py
View File

@@ -1,121 +1,121 @@
from threading import Thread
from typing import Iterator, Optional, Callable
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, TextIteratorStreamer
import torch
from . import LlmEngine
from .. import util
class LocalLlmEngine(LlmEngine):
def __init__(
self,
model_path: str,
temperature: float,
token_limit: int,
api_token: Optional[str],
):
"""
Initialize the LLM Engine with a model path.
Args:
model_path: Path to the model weights to be used.
temperature: Temperature for sampling
api_token: Huggingface API key
token_limit: Maximum number of tokens to generate
"""
self._temperature = temperature
self._token_limit = token_limit
self._tokenizer = AutoTokenizer.from_pretrained(model_path, token=api_token)
model = AutoModelForCausalLM.from_pretrained(
model_path,
return_dict=True,
low_cpu_mem_usage=True,
torch_dtype=torch.bfloat16,
device_map="auto",
trust_remote_code=True,
token=api_token,
)
if self._tokenizer.pad_token_id is None:
self._tokenizer.pad_token_id = self._tokenizer.eos_token_id
if model.config.pad_token_id is None:
model.config.pad_token_id = model.config.eos_token_id
self._pipeline = pipeline(
"text-generation",
model=model,
tokenizer=self._tokenizer,
torch_dtype=torch.bfloat16,
device_map="auto",
return_full_text=False,
)
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.
"""
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": main_context}
]
prompt = self._tokenizer.apply_chat_template(
messages, tokenize=False, add_generation_prompt=True
)
streamer = TextIteratorStreamer(
self._tokenizer,
skip_prompt=True
)
generation_thread = Thread(target=self._pipeline, kwargs=dict(
text_inputs=prompt,
do_sample=True,
temperature=self._temperature,
max_new_tokens=self.token_limit(),
streamer=streamer
))
generation_thread.start()
for text in util.stop_before_value(streamer, '<|eot_id|>'):
yield text
if should_stop():
break
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
"""
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": main_context}
]
prompt = self._tokenizer.apply_chat_template(
messages, tokenize=False, add_generation_prompt=True
)
return len(self._tokenizer.encode(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
else:
return self._pipeline.model.config.max_position_embeddings
from threading import Thread
from typing import Iterator, Optional, Callable
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, TextIteratorStreamer
import torch
from . import LlmEngine
from .. import util
class LocalLlmEngine(LlmEngine):
def __init__(
self,
model_path: str,
temperature: float,
token_limit: int,
api_token: Optional[str],
):
"""
Initialize the LLM Engine with a model path.
Args:
model_path: Path to the model weights to be used.
temperature: Temperature for sampling
token_limit: Maximum number of tokens to generate
api_token: Huggingface API key
"""
self._temperature = temperature
self._token_limit = token_limit
self._tokenizer = AutoTokenizer.from_pretrained(model_path, token=api_token)
model = AutoModelForCausalLM.from_pretrained(
model_path,
return_dict=True,
low_cpu_mem_usage=True,
torch_dtype=torch.bfloat16,
device_map="auto",
trust_remote_code=True,
token=api_token,
)
if self._tokenizer.pad_token_id is None:
self._tokenizer.pad_token_id = self._tokenizer.eos_token_id
if model.config.pad_token_id is None:
model.config.pad_token_id = model.config.eos_token_id
self._pipeline = pipeline(
"text-generation",
model=model,
tokenizer=self._tokenizer,
torch_dtype=torch.bfloat16,
device_map="auto",
return_full_text=False,
)
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.
"""
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": main_context}
]
prompt = self._tokenizer.apply_chat_template(
messages, tokenize=False, add_generation_prompt=True
)
streamer = TextIteratorStreamer(
self._tokenizer,
skip_prompt=True
)
generation_thread = Thread(target=self._pipeline, kwargs=dict(
text_inputs=prompt,
do_sample=True,
temperature=self._temperature,
max_new_tokens=self.token_limit(),
streamer=streamer
))
generation_thread.start()
for text in util.stop_before_value(streamer, '<|eot_id|>'):
yield text
if should_stop():
break
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
"""
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": main_context}
]
prompt = self._tokenizer.apply_chat_template(
messages, tokenize=False, add_generation_prompt=True
)
return len(self._tokenizer.encode(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
else:
return self._pipeline.model.config.max_position_embeddings

427
sia/llm_engine/qwq_llm_engine.py
View File

@@ -1,303 +1,124 @@
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
from pathlib import Path
from threading import Thread
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, TextIteratorStreamer, BitsAndBytesConfig
from typing import Callable, Iterator
import torch
from . import LlmEngine
from .. import util
class QwQLlmEngine(LlmEngine):
def __init__(
self,
model_path: Path,
temperature: float,
token_limit: int = None,
):
"""
Initialize the QwQ LLM Engine.
Args:
model_path: Local path to the model
temperature: Sampling temperature
token_limit: Maximum tokens to generate
"""
self._temperature = temperature
self._token_limit = token_limit
quantization_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_compute_dtype=torch.bfloat16,
bnb_4bit_quant_type="nf4",
bnb_4bit_use_double_quant=True
)
model = AutoModelForCausalLM.from_pretrained(
model_path,
return_dict=True,
device_map="auto",
attn_implementation="flash_attention_2",
use_cache=True,
quantization_config=quantization_config,
)
self._tokenizer = AutoTokenizer.from_pretrained(model_path)
self._pipline = pipeline(
"text-generation",
model=model,
tokenizer=self._tokenizer,
return_full_text=False,
)
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.
"""
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": main_context}
]
text = self._tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=True,
)
streamer = TextIteratorStreamer(
self._tokenizer,
skip_prompt=True,
)
generation_thread = Thread(
target=self._pipline,
kwargs=dict(
text_inputs=text,
do_sample=True,
temperature=self._temperature,
max_new_tokens=self._token_limit,
streamer=streamer,
)
)
generation_thread.start()
for text in util.stop_before_value(streamer, '<|eot_id|>'):
yield text
if should_stop():
break
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
"""
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": main_context}
]
prompt = self._tokenizer.apply_chat_template(
messages, tokenize=False, add_generation_prompt=True
)
return len(self._tokenizer.encode(prompt))
def token_limit(self) -> int:
return self._token_limit