From 6f3d414d179b02c2cb61e0f122600714e8aaeb49 Mon Sep 17 00:00:00 2001 From: Niels Geens Date: Tue, 1 Apr 2025 10:41:24 +0200 Subject: [PATCH] Converted QwQ notebooks to .py files --- sia/__main__.py | 239 ++++++++-------- sia/llm_engine/local_llm_engine.py | 242 ++++++++-------- sia/llm_engine/qwq_llm_engine.py | 427 +++++++++-------------------- tools/train/train/qwq.py | 16 +- tools/train/train/qwq_train.ipynb | 2 +- 5 files changed, 368 insertions(+), 558 deletions(-) diff --git a/sia/__main__.py b/sia/__main__.py index bbc47e6..eb6491e 100644 --- a/sia/__main__.py +++ b/sia/__main__.py @@ -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 \ No newline at end of file diff --git a/sia/llm_engine/local_llm_engine.py b/sia/llm_engine/local_llm_engine.py index 126ab73..887a000 100644 --- a/sia/llm_engine/local_llm_engine.py +++ b/sia/llm_engine/local_llm_engine.py @@ -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 diff --git a/sia/llm_engine/qwq_llm_engine.py b/sia/llm_engine/qwq_llm_engine.py index ebde720..dbb9410 100644 --- a/sia/llm_engine/qwq_llm_engine.py +++ b/sia/llm_engine/qwq_llm_engine.py @@ -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 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 \ No newline at end of file +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 \ No newline at end of file diff --git a/tools/train/train/qwq.py b/tools/train/train/qwq.py index b4e9bfd..7db4086 100644 --- a/tools/train/train/qwq.py +++ b/tools/train/train/qwq.py @@ -6,7 +6,6 @@ Fine-tuning for QwQ model # Unsloth should be imported before transformers to ensure all optimizations are applied. from unsloth import FastLanguageModel, is_bfloat16_supported -from .dataset import Dataset from dataclasses import dataclass from pathlib import Path from transformers import TrainingArguments @@ -15,6 +14,8 @@ from typing import Optional, List import argparse import os +from .dataset import Dataset + @dataclass class Args: def __init__(self, args: Optional[List[str]]): @@ -78,7 +79,7 @@ def main(): load_in_4bit = True, # False for LoRA 16bit fast_inference = True, # Enable vLLM fast inference max_lora_rank = lora_rank, - gpu_memory_utilization = 0.85, # Reduce if out of memory + gpu_memory_utilization = 0.5, # Reduce if out of memory ) model = FastLanguageModel.get_peft_model( @@ -97,12 +98,6 @@ def main(): loftq_config = None, # And LoftQ ) - response_template = tokenizer.apply_chat_template( - [{"role": "assistant", "content": ""}], - tokenize=False, - add_generation_prompt=True - ) - training_args = TrainingArguments( output_dir=str(args.output_dir), num_train_epochs=3, @@ -129,10 +124,6 @@ def main(): train_dataset=dataset.to_transformers_dataset(tokenizer), dataset_text_field="messages", max_seq_length=max_seq_length, - data_collator=DataCollatorForCompletionOnlyLM( - response_template=response_template, - tokenizer=tokenizer - ), ) trainer.train() @@ -140,7 +131,6 @@ def main(): model.save_pretrained_merged( str(args.output_dir), tokenizer=tokenizer, - save_method="merged_16bit" ) if __name__ == "__main__": diff --git a/tools/train/train/qwq_train.ipynb b/tools/train/train/qwq_train.ipynb index ad0156b..085bc87 100644 --- a/tools/train/train/qwq_train.ipynb +++ b/tools/train/train/qwq_train.ipynb @@ -22,7 +22,7 @@ }, { "cell_type": "code", - "execution_count": 2, + "execution_count": null, "metadata": {}, "outputs": [], "source": [