Added tool for finetuning mistral model based on training config
This commit is contained in:
Niels Geens
@@ -1,134 +0,0 @@
|
||||
import json
|
||||
import hashlib
|
||||
import xml.etree.ElementTree as ET
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
class FinetuneDatasetCreator:
|
||||
"""Creates JSONL finetune dataset from iteration XML files"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
iterations_dir: Path,
|
||||
system_prompt_file: Path,
|
||||
action_schema_file: Path,
|
||||
output_file: Path
|
||||
):
|
||||
"""
|
||||
Initialize the dataset creator
|
||||
|
||||
Args:
|
||||
iterations_dir: Directory containing iteration XML files
|
||||
system_prompt_file: Path to system prompt file
|
||||
action_schema_file: Path to action schema file
|
||||
output_file: Path where JSONL dataset will be written
|
||||
"""
|
||||
self.iterations_dir = Path(iterations_dir)
|
||||
self.system_prompt_file = Path(system_prompt_file)
|
||||
self.action_schema_file = Path(action_schema_file)
|
||||
self.output_file = Path(output_file)
|
||||
|
||||
# Read and hash system prompt and 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 _calculate_hash(self, content: str) -> str:
|
||||
"""Calculate SHA256 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 messages dictionary
|
||||
|
||||
Returns None if hashes don't match or parsing fails
|
||||
"""
|
||||
try:
|
||||
tree = ET.parse(file_path)
|
||||
root = tree.getroot()
|
||||
|
||||
# Verify hashes
|
||||
if root.get('system_prompt_hash') != self.system_prompt_hash:
|
||||
print(f"System prompt hash mismatch in {file_path}")
|
||||
if root.get('action_schema_hash') != self.action_schema_hash:
|
||||
print(f"Action schema hash mismatch in {file_path}")
|
||||
|
||||
# Get context and response
|
||||
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
|
||||
|
||||
# Create messages list
|
||||
messages = [
|
||||
{
|
||||
"role": "system",
|
||||
"content": self.system_prompt + self.action_schema
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": context
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": response
|
||||
}
|
||||
]
|
||||
|
||||
return {"messages": messages}
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error processing {file_path}: {str(e)}")
|
||||
return None
|
||||
|
||||
def create_dataset(self) -> int:
|
||||
"""
|
||||
Create JSONL dataset from all valid iteration files
|
||||
|
||||
Returns:
|
||||
Number of samples written to dataset
|
||||
"""
|
||||
sample_count = 0
|
||||
|
||||
# Create output directory if needed
|
||||
self.output_file.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
with open(self.output_file, 'w', encoding='utf-8') as f:
|
||||
# Process each XML file in iterations directory
|
||||
for xml_file in sorted(self.iterations_dir.rglob('*.xml')):
|
||||
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 main():
|
||||
"""Command line interface"""
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser(description='Create finetune dataset from iteration XML files')
|
||||
parser.add_argument('iterations_dir', type=str, help='Directory containing iteration XML files')
|
||||
parser.add_argument('system_prompt_file', type=str, help='Path to system prompt file')
|
||||
parser.add_argument('action_schema_file', type=str, help='Path to action schema file')
|
||||
parser.add_argument('output_file', type=str, help='Path for output JSONL dataset')
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
creator = FinetuneDatasetCreator(
|
||||
iterations_dir=args.iterations_dir,
|
||||
system_prompt_file=args.system_prompt_file,
|
||||
action_schema_file=args.action_schema_file,
|
||||
output_file=args.output_file
|
||||
)
|
||||
|
||||
creator.create_dataset()
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -299,21 +299,13 @@ The training configuration is defined in `/training/config.yaml`, which specifie
|
||||
model:
|
||||
system_prompt_path: "system_prompt.md"
|
||||
action_schema: "action_schema.xsd"
|
||||
|
||||
training_data:
|
||||
conversations:
|
||||
- path: "training/data/general/basic_interactions/"
|
||||
description: "General user interactions"
|
||||
hash: "abc123"
|
||||
tasks:
|
||||
- path: "training/data/code_generation/"
|
||||
description: "Code writing examples"
|
||||
hash: "def456"
|
||||
|
||||
training_params:
|
||||
batch_size: 32
|
||||
params:
|
||||
learning_rate: 1e-5
|
||||
epochs: 3
|
||||
data:
|
||||
- "training/clean_start/"
|
||||
- "training/delete_indicated_entries/"
|
||||
- "training/list_entries_to_delete/"
|
||||
```
|
||||
|
||||
## Continuous Operation
|
||||
@@ -342,3 +334,13 @@ This requires:
|
||||
- Monitor system resource usage
|
||||
- Prevent training impact on user tasks
|
||||
- Clean up old data regularly
|
||||
|
||||
# TODO
|
||||
|
||||
- Fix training config
|
||||
- Write training script
|
||||
- implement stdio + auto mode
|
||||
- Write setup script
|
||||
- Explain challenge report card
|
||||
- Document report card tool
|
||||
- Write report card tool
|
||||
260
tools/train/train_mistral.py
Normal file
260
tools/train/train_mistral.py
Normal file
@@ -0,0 +1,260 @@
|
||||
#!/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())
|
||||
Reference in New Issue
Block a user