WIP Simplify training

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"

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")

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
@@ -79,256 +54,11 @@ class Config:
     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())
+    main()

tools/train/train/util.py

@@ -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:
-    # <s>[INST] {system + user content} [/INST] {assistant response} </s>
-    
-    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"<s>[INST] {instruction} [/INST] {assistant_content} </s>"
-
-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}")

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/"

web/.gitignore

@@ -1,4 +1,5 @@
-node_modules
-dist
-coverage
 .DS_Store
+coverage
+dist
+node_modules
+package-lock.json