From 3594c8150aca4d1f49aa417a27dbcdc69f3bdbb0 Mon Sep 17 00:00:00 2001 From: Niels Geens Date: Mon, 24 Mar 2025 10:53:51 +0000 Subject: [PATCH] WIP Simplify training --- scripts/bootstrap.sh | 0 scripts/collect.sh | 0 scripts/deploy.sh | 0 scripts/exec.sh | 0 scripts/gitea_keys.sh | 0 scripts/github_keys.sh | 0 scripts/install.sh | 0 scripts/local.sh | 0 scripts/restart.sh | 0 scripts/test.sh | 0 tools/train/train/__init__.py | 8 - tools/train/train/dataset.py | 132 ++++++++++++++++ tools/train/train/qwq.py | 290 ++-------------------------------- tools/train/train/util.py | 195 ----------------------- training/config.yaml | 3 - web/.gitignore | 7 +- 16 files changed, 146 insertions(+), 489 deletions(-) mode change 100644 => 100755 scripts/bootstrap.sh mode change 100644 => 100755 scripts/collect.sh mode change 100644 => 100755 scripts/deploy.sh mode change 100644 => 100755 scripts/exec.sh mode change 100644 => 100755 scripts/gitea_keys.sh mode change 100644 => 100755 scripts/github_keys.sh mode change 100644 => 100755 scripts/install.sh mode change 100644 => 100755 scripts/local.sh mode change 100644 => 100755 scripts/restart.sh mode change 100644 => 100755 scripts/test.sh create mode 100644 tools/train/train/dataset.py delete mode 100644 tools/train/train/util.py diff --git a/scripts/bootstrap.sh b/scripts/bootstrap.sh old mode 100644 new mode 100755 diff --git a/scripts/collect.sh b/scripts/collect.sh old mode 100644 new mode 100755 diff --git a/scripts/deploy.sh b/scripts/deploy.sh old mode 100644 new mode 100755 diff --git a/scripts/exec.sh b/scripts/exec.sh old mode 100644 new mode 100755 diff --git a/scripts/gitea_keys.sh b/scripts/gitea_keys.sh old mode 100644 new mode 100755 diff --git a/scripts/github_keys.sh b/scripts/github_keys.sh old mode 100644 new mode 100755 diff --git a/scripts/install.sh b/scripts/install.sh old mode 100644 new mode 100755 diff --git a/scripts/local.sh b/scripts/local.sh old mode 100644 new mode 100755 diff --git a/scripts/restart.sh b/scripts/restart.sh old mode 100644 new mode 100755 diff --git a/scripts/test.sh b/scripts/test.sh old mode 100644 new mode 100755 diff --git a/tools/train/train/__init__.py b/tools/train/train/__init__.py index b2a0d92..e69de29 100644 --- a/tools/train/train/__init__.py +++ b/tools/train/train/__init__.py @@ -1,8 +0,0 @@ -""" -SIA Training Tool - -This package provides utilities for fine-tuning language models used by SIA. -Supports DeepSeek and Mistral models. -""" - -__version__ = "0.1.0" \ No newline at end of file diff --git a/tools/train/train/dataset.py b/tools/train/train/dataset.py new file mode 100644 index 0000000..1f6bbbc --- /dev/null +++ b/tools/train/train/dataset.py @@ -0,0 +1,132 @@ +from dataclasses import dataclass +from pathlib import Path +from typing import Dict, List, Optional, Tuple, Any, Iterator +import hashlib +import json +import yaml +import xml.etree.ElementTree as ET + + +class Dataset: + """Training dataset from XML iteration files""" + + def __init__(self, config_filename: str): + with open(config_filename) as f: + config_data = yaml.safe_load(f) + + data_paths = [Path(p) for p in config_data['data']] + self.files = self._find_xml_files(data_paths) + + self.system_prompt_file = Path(config_data['model']['system_prompt_path']) + self.action_schema_file = Path(config_data['model']['action_schema']) + + self.system_prompt = self.system_prompt_file.read_text() + self.system_prompt_hash = self._calculate_hash(self.system_prompt) + + self.action_schema = self.action_schema_file.read_text() + self.action_schema_hash = self._calculate_hash(self.action_schema) + + def _find_xml_files(self, data_paths: List[Path]) -> List[Path]: + """Find all XML files in the given data paths""" + xml_files = list() + for path in data_paths: + if not path.exists(): + raise Exception(f"Data path not found: {path}") + xml_files.extend(path.rglob('*.xml')) + return xml_files + + def _calculate_hash(self, content: str) -> str: + """Calculate SHA-256 hash of content""" + return hashlib.sha256(content.encode()).hexdigest() + + def _parse_iteration_file(self, file_path: Path) -> Dict: + """Parse a single iteration XML file into a training example""" + tree = ET.parse(file_path) + root = tree.getroot() + + context_elem = root.find('context') + response_elem = root.find('response') + + context = context_elem.text + response = response_elem.text + + return { + "messages": [ + { + "role": "system", + "content": self.system_prompt + "\n" + self.action_schema + }, + { + "role": "user", + "content": context + }, + { + "role": "assistant", + "content": response + } + ] + } + + def __len__(self) -> int: + """Return the number of samples in the dataset""" + return len(self.files) + + def __getitem__(self, idx: int) -> Dict: + """Indexing for a single sample""" + if idx < 0 or idx >= len(self): + raise IndexError(f"Index {idx} out of range for dataset with {len(self)} samples") + file_path = self.files[idx] + return self._parse_iteration_file(file_path) + + def __iter__(self) -> Iterator[Dict]: + """Allow iteration over samples""" + for i in range(len(self)): + yield self[i] + + def to_list(self) -> List[Dict]: + """Convert dataset to a list""" + results = [] + for i in range(len(self)): + results.append(self[i]) + return results + + def validate(self) -> None: + """Validate XML files""" + print(f"Validating {len(self.files)} XML files...") + + for i in range(len(self.files)): + self.validate_sample(i) + + print(f"Validation complete. Found {len(self.files)} valid files.") + + def validate_sample(self, index: int) -> None: + file = self.files[index] + print("file:", file) + tree = ET.parse(file) + root = tree.getroot() + + # Check system prompt hash + file_system_hash = root.get('system_prompt_hash') + if file_system_hash != self.system_prompt_hash: + print(f"WARNING: System prompt hash mismatch in {file_path}") + + # Check action schema hash + file_schema_hash = root.get('action_schema_hash') + if file_schema_hash != self.action_schema_hash: + print(f"WARNING: Action schema hash mismatch in {file_path}") + + # Check for required elements + context_elem = root.find('context') + response_elem = root.find('response') + + if context_elem is None: + raise Exception(f"Missing context element") + + if response_elem is None: + raise Exception(f"Missing response element") + + if not context_elem.text: + raise Exception(f"Empty context") + + if not response_elem.text: + raise Exception(f"Empty response") \ No newline at end of file diff --git a/tools/train/train/qwq.py b/tools/train/train/qwq.py index ca0329b..ea685ab 100644 --- a/tools/train/train/qwq.py +++ b/tools/train/train/qwq.py @@ -1,27 +1,15 @@ #!/root/venvs/train/bin/python """ -Fine-tuning script for QwQ models to support SIA's action schema. -Supports both full and LoRA finetuning methods. +Fine-tuning for QwQ model """ -import argparse -import os -import sys -import torch +from .dataset import Dataset from dataclasses import dataclass +import argparse from pathlib import Path -import json -import logging -import gc - -# Set up logging -logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') -logger = logging.getLogger(__name__) - -# Import from shared library -from .util import prepare_training_data +import os @dataclass -class Config: +class Args: def __init__(self): parser = argparse.ArgumentParser(description='Train SIA model using QwQ') parser.add_argument( @@ -48,19 +36,6 @@ class Config: default=os.environ.get('SIA_HF_API_KEY'), help='HuggingFace API key' ) - parser.add_argument( - '--method', - type=str, - choices=['lora', 'qlora', 'full'], - default='qlora', - help='Finetuning method: LoRA, QLoRA (quantized LoRA), or full-model' - ) - parser.add_argument( - '--device', - type=str, - default='auto', - help='Override device (cpu, cuda, auto) from config' - ) self.args = parser.parse_args() @property @@ -78,257 +53,12 @@ class Config: @property def api_key(self) -> str: return self.args.api_key - - @property - def device(self) -> str: - return self.args.device - - @property - def method(self) -> str: - return self.args.method - -def format_data_for_qwq(training_data): - """ - Format training data for QwQ model focusing on action schema formats. - Ensures each example shows the model how to directly use action elements. - """ - formatted_data = [] - - for sample in training_data: - # Get the system prompt, context, and response - system_content = "" - context_content = "" - response_content = "" - - for message in sample.get("messages", []): - if message["role"] == "system": - system_content = message["content"] - elif message["role"] == "user": - context_content = message["content"] - elif message["role"] == "assistant": - response_content = message["content"] - - # Create conversations with explicit instruction to use action schema - formatted_data.append({ - "conversations": [ - {"role": "system", "content": system_content}, - {"role": "user", "content": context_content}, - {"role": "assistant", "content": response_content} - ] - }) - - logger.info(f"Formatted {len(formatted_data)} examples for QwQ training") - return formatted_data - -def train_model_lora(config, training_data, train_params): - """ - Train QwQ model using LoRA or QLoRA for parameter-efficient fine-tuning. - This is the recommended approach for most use cases. - """ - try: - # Import required libraries - from transformers import ( - AutoModelForCausalLM, AutoTokenizer, - TrainingArguments, DataCollatorForSeq2Seq - ) - from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training - from datasets import Dataset - from trl import SFTTrainer - - except ImportError as e: - logger.error(f"Error importing required libraries: {e}") - logger.error("Please ensure transformers, peft, and trl are installed.") - sys.exit(1) - - # Format data specifically for QwQ - formatted_data = format_data_for_qwq(training_data) - dataset = Dataset.from_list(formatted_data) - - logger.info(f"Starting QwQ fine-tuning using {config.method}") - logger.info(f"Base model: {config.base_model}") - logger.info(f"Device: {config.device}") - - # Configure device mapping and precision - if torch.cuda.is_available() and config.device != "cpu": - logger.info("Using GPU for training") - device_map = "auto" - - # Configure precision based on method - if config.method == "qlora": - dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16 - load_in_4bit = True - load_in_8bit = False - else: - dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16 - load_in_4bit = False - load_in_8bit = False - else: - logger.info("Using CPU for training") - device_map = "cpu" - dtype = torch.float32 - load_in_4bit = False - load_in_8bit = False - - # Configure quantization for QLoRA - if config.method == "qlora": - from transformers import BitsAndBytesConfig - logger.info("Setting up 4-bit quantization for QLoRA") - compute_dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16 - - bnb_config = BitsAndBytesConfig( - load_in_4bit=True, - bnb_4bit_quant_type="nf4", - bnb_4bit_compute_dtype=compute_dtype, - bnb_4bit_use_double_quant=True - ) - else: - bnb_config = None - - # Load tokenizer - logger.info(f"Loading tokenizer from {config.base_model}") - tokenizer = AutoTokenizer.from_pretrained( - config.base_model, - token=config.api_key, - trust_remote_code=True - ) - - if tokenizer.pad_token is None: - tokenizer.pad_token = tokenizer.eos_token - - # Load model - logger.info(f"Loading model from {config.base_model}") - model = AutoModelForCausalLM.from_pretrained( - config.base_model, - torch_dtype=dtype, - device_map=device_map, - quantization_config=bnb_config, - token=config.api_key, - trust_remote_code=True - ) - - # Configure LoRA - if config.method in ["lora", "qlora"]: - if config.method == "qlora": - model = prepare_model_for_kbit_training(model) - - logger.info("Setting up LoRA configuration") - target_modules = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"] - - lora_config = LoraConfig( - r=16, - lora_alpha=32, - target_modules=target_modules, - lora_dropout=0.05, - bias="none", - task_type="CAUSAL_LM" - ) - - model = get_peft_model(model, lora_config) - - # Create output directory - config.output_dir.mkdir(parents=True, exist_ok=True) - - # Configure training arguments - batch_size = train_params.per_device_batch_size - gradient_accumulation = train_params.gradient_accumulation_steps - - # Scale down batch size based on model - if "32B" in config.base_model and batch_size > 1: - batch_size = 1 - gradient_accumulation *= 2 - - training_args = TrainingArguments( - output_dir=str(config.output_dir), - per_device_train_batch_size=batch_size, - gradient_accumulation_steps=gradient_accumulation, - learning_rate=train_params.learning_rate, - num_train_epochs=train_params.epochs, - logging_steps=10, - save_strategy="epoch", - save_total_limit=2, - fp16=dtype == torch.float16, - bf16=dtype == torch.bfloat16, - report_to="none", - remove_unused_columns=False, - optim="adamw_torch", - weight_decay=0.01, - max_grad_norm=0.3, - warmup_ratio=0.03, - lr_scheduler_type="cosine", - seed=42 - ) - - # Set up trainer - logger.info("Setting up trainer") - trainer = SFTTrainer( - model=model, - args=training_args, - train_dataset=dataset, - tokenizer=tokenizer, - max_seq_length=train_params.max_seq_length, - dataset_text_field="conversations", - packing=False - ) - - # Start training - logger.info("Starting training") - trainer.train() - - # Save the final model - logger.info(f"Saving model to {config.output_dir}") - trainer.save_model(config.output_dir) - tokenizer.save_pretrained(config.output_dir) - - # Create metadata file - with open(config.output_dir / "training_info.json", "w") as f: - json.dump({ - "base_model": config.base_model, - "method": config.method, - "learning_rate": train_params.learning_rate, - "epochs": train_params.epochs, - "dataset_size": len(dataset), - "batch_size": batch_size, - "gradient_accumulation": gradient_accumulation - }, f, indent=2) - - logger.info("Training complete!") - return True def main(): - # Initialize configuration - config = Config() - - # Prepare training data from config - training_data, train_params = prepare_training_data(config.config_path) - - if not training_data: - logger.error("No valid training data found. Exiting.") - return 1 - - # Force garbage collection - gc.collect() - - # Train using appropriate method - if config.method in ["lora", "qlora"]: - success = train_model_lora(config, training_data, train_params) - else: - logger.error(f"Training method '{config.method}' not yet implemented") - return 1 - - if success: - # Create symlink to current - current_link = Path("/root/models/current") - if os.path.exists(current_link) or os.path.islink(current_link): - os.unlink(current_link) - os.symlink(config.output_dir, current_link, target_is_directory=True) - - logger.info(f"Training complete. Model saved to {config.output_dir}") - logger.info(f"Symlink created at {current_link}") - - return 0 - else: - logger.error("Training failed") - return 1 + args = Args() + dataset = Dataset(args.config_path) + dataset.validate() + print(dataset[3]) if __name__ == "__main__": - sys.exit(main()) \ No newline at end of file + main() \ No newline at end of file diff --git a/tools/train/train/util.py b/tools/train/train/util.py deleted file mode 100644 index b1ca6c0..0000000 --- a/tools/train/train/util.py +++ /dev/null @@ -1,195 +0,0 @@ -""" -Shared library for SIA model training functionality. -Contains common code for both API-based and local training. -""" -from dataclasses import dataclass -from pathlib import Path -from typing import Dict, List, Optional, Set, Tuple, Any -import hashlib -import json -import subprocess -import sys -import xml.etree.ElementTree as ET -import yaml - -@dataclass -class TrainingParams: - """Parameters for model training""" - learning_rate: float - epochs: int - batch_size: int - max_seq_length: int - quantization: str - per_device_batch_size: int - gradient_accumulation_steps: int - mixed_precision: str - - @classmethod - def from_dict(cls, config_dict: Dict[str, Any]) -> 'TrainingParams': - """Create from config dictionary with defaults""" - return cls( - learning_rate=float(config_dict.get('learning_rate', 1e-5)), - epochs=int(config_dict.get('epochs', 1)), - batch_size=int(config_dict.get('batch_size', 1)), - max_seq_length=int(config_dict.get('max_seq_length', 1024)), - quantization=config_dict.get('quantization', '4bit'), - per_device_batch_size=int(config_dict.get('per_device_batch_size', 1)), - gradient_accumulation_steps=int(config_dict.get('gradient_accumulation_steps', 8)), - mixed_precision=config_dict.get('mixed_precision', 'no') - ) - -class DatasetCreator: - """Creates training datasets from XML iteration files""" - - def __init__( - self, - xml_files: Set[Path], - system_prompt_file: Path, - action_schema_file: Path - ): - self.xml_files = xml_files - self.system_prompt_file = Path(system_prompt_file) - self.action_schema_file = Path(action_schema_file) - - self.system_prompt = self.system_prompt_file.read_text() - self.system_prompt_hash = self._calculate_hash(self.system_prompt) - - self.action_schema = self.action_schema_file.read_text() - self.action_schema_hash = self._calculate_hash(self.action_schema) - - def _calculate_hash(self, content: str) -> str: - """Calculate SHA-256 hash of content""" - return hashlib.sha256(content.encode()).hexdigest() - - def _parse_iteration_file(self, file_path: Path) -> Optional[Dict]: - """Parse a single iteration XML file into a training example""" - try: - tree = ET.parse(file_path) - root = tree.getroot() - - # Check hashes to ensure compatibility - if root.get('system_prompt_hash') != self.system_prompt_hash: - print(f"System prompt hash mismatch in {file_path}") - return None - if root.get('action_schema_hash') != self.action_schema_hash: - print(f"Action schema hash mismatch in {file_path}") - return None - - context_elem = root.find('context') - response_elem = root.find('response') - - if context_elem is None or response_elem is None: - print(f"Missing context or response elements in {file_path}") - return None - - context = context_elem.text - response = response_elem.text - - if not context or not response: - print(f"Empty context or response in {file_path}") - return None - - return { - "messages": [ - { - "role": "system", - "content": self.system_prompt + "\n" + self.action_schema - }, - { - "role": "user", - "content": context - }, - { - "role": "assistant", - "content": response - } - ] - } - - except Exception as e: - print(f"Error processing {file_path}: {str(e)}") - return None - - def create_dataset(self) -> List[Dict]: - """Create a dataset from all valid XML files""" - samples = [] - total_files = len(self.xml_files) - print(f"Processing {total_files} XML files...") - - for i, xml_file in enumerate(sorted(self.xml_files)): - if i % 10 == 0: - print(f"Processed {i}/{total_files} files...") - - sample = self._parse_iteration_file(xml_file) - if sample: - samples.append(sample) - - print(f"Created dataset with {len(samples)} samples from {total_files} files") - return samples - -def find_xml_files(data_paths: List[Path]) -> Set[Path]: - """Find all XML files in the given data paths""" - xml_files = set() - for path in data_paths: - if not path.exists(): - print(f"Error: Data path not found: {path}") - sys.exit(1) - xml_files.update(path.rglob('*.xml')) - return xml_files - -def format_chat_for_mistral(messages): - """Format messages for Mistral chat format""" - # Mistral uses a specific chat format: - # [INST] {system + user content} [/INST] {assistant response} - - system_content = "" - user_content = "" - assistant_content = "" - - for msg in messages: - role = msg["role"] - content = msg["content"] - - if role == "system": - system_content = content - elif role == "user": - user_content = content - elif role == "assistant": - assistant_content = content - - # Combine system and user content for the instruction - instruction = system_content - if instruction and user_content: - instruction += "\n\n" - instruction += user_content - - # Format according to Mistral chat template - return f"[INST] {instruction} [/INST] {assistant_content} " - -def prepare_training_data(config_path: Path) -> Tuple[List[Dict], TrainingParams]: - """Prepare training data from config and XML files""" - with open(config_path) as f: - config_data = yaml.safe_load(f) - - data_paths = [Path(p) for p in config_data['data']] - xml_files = find_xml_files(data_paths) - - creator = DatasetCreator( - xml_files=xml_files, - system_prompt_file=config_data['model']['system_prompt_path'], - action_schema_file=config_data['model']['action_schema'] - ) - - training_data = creator.create_dataset() - - train_params = TrainingParams.from_dict(config_data['params']) - - return training_data, train_params - -def save_jsonl_dataset(data: List[Dict], output_path: Path) -> None: - """Save dataset in JSONL format""" - with open(output_path, 'w', encoding='utf-8') as f: - for sample in data: - json.dump(sample, f, ensure_ascii=False) - f.write('\n') - print(f"Saved dataset with {len(data)} samples to {output_path}") diff --git a/training/config.yaml b/training/config.yaml index 599af77..baace4f 100644 --- a/training/config.yaml +++ b/training/config.yaml @@ -1,9 +1,6 @@ model: system_prompt_path: "/root/sia/system_prompt.md" action_schema: "/root/sia/action_schema.xsd" -params: - learning_rate: 1e-5 - epochs: 3 data: - "/root/sia/training/clean_start/" - "/root/sia/training/delete_indicated_entries/" diff --git a/web/.gitignore b/web/.gitignore index 8a771b4..bfe46fb 100644 --- a/web/.gitignore +++ b/web/.gitignore @@ -1,4 +1,5 @@ -node_modules -dist +.DS_Store coverage -.DS_Store \ No newline at end of file +dist +node_modules +package-lock.json \ No newline at end of file