#!/usr/bin/env python3 from dataclasses import dataclass from datetime import datetime from dotenv import load_dotenv from pathlib import Path from typing import Dict, List, Optional, Set import argparse import hashlib import json import os import requests import subprocess import sys import tempfile import xml.etree.ElementTree as ET import yaml @dataclass class Config: def __init__(self): load_dotenv() parser = argparse.ArgumentParser(description='Train SIA model using Mistral API') parser.add_argument( '--config', type=Path, default=os.getenv('SIA_TRAINING_CONFIG', '/root/sia/training/config.yaml'), help='Path to config file' ) parser.add_argument( '--model', type=str, default=os.getenv('SIA_MISTRAL_MODEL', 'mistral-large-latest'), help='Base model for fine-tuning' ) parser.add_argument( '--api-key', type=str, default=os.getenv('SIA_MISTRAL_API_KEY'), help='Mistral API key' ) self.args = parser.parse_args() @property def config_path(self) -> Path: return self.args.config @property def model(self) -> str: return self.args.model @property def api_key(self) -> str: return self.args.api_key class FinetuneDatasetCreator: def __init__( self, xml_files: Set[Path], system_prompt_file: Path, action_schema_file: Path, output_file: Path ): self.xml_files = xml_files self.system_prompt_file = Path(system_prompt_file) self.action_schema_file = Path(action_schema_file) self.output_file = Path(output_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: return hashlib.sha256(content.encode()).hexdigest() def _parse_iteration_file(self, file_path: Path) -> Optional[Dict]: try: tree = ET.parse(file_path) root = tree.getroot() 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 = root.find('context').text response = root.find('response').text if not context or not response: print(f"Missing context or response in {file_path}") return None return { "messages": [ { "role": "system", "content": self.system_prompt + 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) -> int: sample_count = 0 self.output_file.parent.mkdir(parents=True, exist_ok=True) with open(self.output_file, 'w', encoding='utf-8') as f: for xml_file in sorted(self.xml_files): sample = self._parse_iteration_file(xml_file) if sample: json.dump(sample, f, ensure_ascii=False) f.write('\n') sample_count += 1 print(f"Created dataset with {sample_count} samples at {self.output_file}") return sample_count def find_xml_files(data_paths: List[Path]) -> Set[Path]: 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 check_git_status(paths: list[Path]) -> str: try: for path in paths: result = subprocess.run(['git', 'status', '--porcelain', str(path)], capture_output=True, text=True) if result.stdout.strip(): print(f"Error: Uncommitted changes in {path}") print(result.stdout) sys.exit(1) result = subprocess.run(['git', 'rev-parse', 'HEAD'], capture_output=True, text=True) return result.stdout.strip() except subprocess.CalledProcessError as e: print(f"Git command failed: {e}") sys.exit(1) def create_combined_dataset(xml_files: Set[Path], config_data: dict, tmp_dir: Path) -> list: tmp_file = tmp_dir / "dataset.jsonl" creator = FinetuneDatasetCreator( xml_files=xml_files, system_prompt_file=config_data['model']['system_prompt_path'], action_schema_file=config_data['model']['action_schema'], output_file=tmp_file ) creator.create_dataset() with open(tmp_file) as f: return [json.loads(line) for line in f] def prepare_training_data(config: Config) -> tuple[list, dict, str]: with open(config.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) paths = list(xml_files) paths.append(config.config_path) paths.append(Path(config_data['model']['system_prompt_path'])) paths.append(Path(config_data['model']['action_schema'])) commit_hash = check_git_status(paths) with tempfile.TemporaryDirectory() as tmp_dir: training_data = create_combined_dataset(xml_files, config_data, Path(tmp_dir)) train_params = { 'learning_rate': config_data['params']['learning_rate'], 'epochs': config_data['params']['epochs'] } return training_data, train_params, commit_hash def upload_file(api_key: str, file_path: Path) -> str: url = "https://api.mistral.ai/v1/files" headers = { "Authorization": f"Bearer {api_key}" } files = { "file": ("dataset.jsonl", open(file_path, "rb"), "application/jsonl"), "purpose": (None, "fine-tune") } response = requests.post(url, headers=headers, files=files) if response.status_code != 200: print(f"Error uploading file: {response.text}") sys.exit(1) return response.json()["id"] def main(): config = Config() if not config.api_key: print("Error: Mistral API key not found. Set SIA_MISTRAL_API_KEY environment variable.") return 1 training_data, train_params, commit_hash = prepare_training_data(config) model_name = f"sia_{commit_hash}" # Create temp file and upload with tempfile.NamedTemporaryFile(mode='w', suffix='.jsonl', delete=False) as f: for sample in training_data: json.dump(sample, f) f.write('\n') try: file_id = upload_file(config.api_key, Path(f.name)) # Create fine-tuning job headers = { "Authorization": f"Bearer {config.api_key}", "Content-Type": "application/json" } data = { "model": config.model, "training_files": [{"file_id": file_id, "weight": 1}], "hyperparameters": train_params } response = requests.post( "https://api.mistral.ai/v1/fine_tuning/jobs", headers=headers, json=data ) if response.status_code != 200: print(f"Error creating fine-tuning job: {response.text}") return 1 job_id = response.json()["id"] print(f"Started fine-tuning job: {model_name}") print(f"Job ID: {job_id}") print(f"Check status: curl -H 'Authorization: Bearer {config.api_key}' https://api.mistral.ai/v1/fine_tuning/jobs/{job_id}") finally: os.unlink(f.name) return 0 if __name__ == "__main__": exit(main())