142 lines
3.9 KiB
Python
142 lines
3.9 KiB
Python
#!/root/venvs/train/bin/python
|
|
"""
|
|
Script for fine-tuning Mistral models for SIA using the Mistral API.
|
|
"""
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
import argparse
|
|
import json
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
import requests
|
|
|
|
# Import from our shared library
|
|
from .util import TrainingParams, DatasetCreator
|
|
|
|
@dataclass
|
|
class Config:
|
|
def __init__(self):
|
|
parser = argparse.ArgumentParser(description='Train SIA model using Mistral API')
|
|
parser.add_argument(
|
|
'--config',
|
|
type=Path,
|
|
default=Path('/root/sia/training/config.yaml'),
|
|
help='Path to config file'
|
|
)
|
|
parser.add_argument(
|
|
'--model',
|
|
type=str,
|
|
default='mistral-large-latest',
|
|
help='Base model for fine-tuning'
|
|
)
|
|
parser.add_argument(
|
|
'--api-key',
|
|
type=str,
|
|
default=os.environ.get('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
|
|
|
|
def upload_file(api_key: str, file_path: Path) -> str:
|
|
"""Upload a file to the Mistral API and return the file ID"""
|
|
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 start_finetune_job(api_key: str, model: str, file_id: str, params: sia_train_lib.TrainingParams):
|
|
"""Start a fine-tuning job on the Mistral API"""
|
|
headers = {
|
|
"Authorization": f"Bearer {api_key}",
|
|
"Content-Type": "application/json"
|
|
}
|
|
data = {
|
|
"model": model,
|
|
"training_files": [{"file_id": file_id, "weight": 1}],
|
|
"hyperparameters": {
|
|
"learning_rate": params.learning_rate,
|
|
"epochs": params.epochs
|
|
}
|
|
}
|
|
|
|
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 None
|
|
|
|
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 = sia_train_lib.prepare_training_data(config.config_path)
|
|
|
|
if not training_data:
|
|
print("No valid training data found. Exiting.")
|
|
return 1
|
|
|
|
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, ensure_ascii=False)
|
|
f.write('\n')
|
|
|
|
try:
|
|
file_id = upload_file(config.api_key, Path(f.name))
|
|
|
|
# Start fine-tuning job
|
|
job_id = start_finetune_job(
|
|
api_key=config.api_key,
|
|
model=config.model,
|
|
file_id=file_id,
|
|
params=train_params
|
|
)
|
|
|
|
if not job_id:
|
|
return 1
|
|
|
|
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())
|