260 lines
8.5 KiB
Python
260 lines
8.5 KiB
Python
#!/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', '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()) |