SIA/tools/train/train/util.py

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