llm/SIA
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

wip deepseek r1

Browse Source
This commit is contained in:
Geens
2025-03-02 22:01:24 +01:00
parent b7e95d7398
commit b64f8d7d33
40 changed files with 6654 additions and 1859 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
tools/itb/bin/itb_click
View File

@@ -1,4 +1,4 @@
#!/usr/bin/env python3
#!/root/venvs/itb/bin/python
import sys
import os
import time

2
tools/itb/bin/itb_cursor
View File

@@ -1,4 +1,4 @@
#!/usr/bin/env python3
#!/root/venvs/itb/bin/python
import sys
import os
import time

2
tools/itb/bin/itb_input
View File

@@ -1,4 +1,4 @@
#!/usr/bin/env python3
#!/root/venvs/itb/bin/python
import sys
import os
import time

2
tools/itb/bin/itb_navigate
View File

@@ -1,4 +1,4 @@
#!/usr/bin/env python3
#!/root/venvs/itb/bin/python
import sys
import os
import time

2
tools/itb/bin/itb_refresh
View File

@@ -1,4 +1,4 @@
#!/usr/bin/env python3
#!/root/venvs/itb/bin/python
import sys
import os
import time

2
tools/itb/bin/itb_screenshot
View File

@@ -1,4 +1,4 @@
#!/usr/bin/env python3
#!/root/venvs/itb/bin/python
import sys
import os
import argparse

2
tools/itb/bin/itb_scroll
View File

@@ -1,4 +1,4 @@
#!/usr/bin/env python3
#!/root/venvs/itb/bin/python
import sys
import os
import time

2
tools/itb/bin/itb_start
View File

@@ -1,4 +1,4 @@
#!/usr/bin/env python3
#!/root/venvs/itb/bin/python
import random
import subprocess
import sys

23
tools/itb/setup.py
View File

@@ -18,14 +18,17 @@ setup(
'selenium>=4.0.0',
'webdriver-manager>=3.8.0',
'click>=8.0.0',
'beautifulsoup4>=4.9.0'
'beautifulsoup4>=4.9.0',
'pytest>=7.0.0',
'pytest-cov>=4.0.0',
'black>=22.0.0',
'flake8>=4.0.0'
],
extras_require={
'dev': [
'pytest>=7.0.0',
'pytest-cov>=4.0.0',
'black>=22.0.0',
'flake8>=4.0.0'
]
}
)
classifiers=[
'Development Status :: 3 - Alpha',
'Intended Audience :: Developers',
'Programming Language :: Python :: 3',
'Programming Language :: Python :: 3.10',
],
python_requires='>=3.10',
)

10
tools/train/bin/train_deepseek Normal file
View File

@@ -0,0 +1,10 @@
#!/root/venvs/train/bin/python
"""
Command-line utility for fine-tuning DeepSeek models using Unsloth.
Always trains from a base model to create a new fine-tuned model.
"""
import sys
from train.unsloth_deepseek import main
if __name__ == "__main__":
sys.exit(main())

9
tools/train/bin/train_mistral Normal file
View File

@@ -0,0 +1,9 @@
#!/root/venvs/train/bin/python
"""
Command-line utility for fine-tuning Mistral models using Mistral API.
"""
import sys
from train.mistral_api import main
if __name__ == "__main__":
sys.exit(main())

68
tools/train/readme.md Normal file
View File

@@ -0,0 +1,68 @@
# SIA Training Tool
This tool provides command-line utilities for fine-tuning SIA's language models.
## Supported Models
- DeepSeek R1 models (including distilled versions)
- Mistral models
## Commands
### train_deepseek
Fine-tune DeepSeek models using Unsloth optimization.
```bash
train_deepseek --base-model deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B --output-dir /root/models/DeepSeek-R1-Distill-Qwen-1.5B
```
Options:
- `--config`: Path to training configuration file (default: /root/sia/training/config.yaml)
- `--base-model`: HuggingFace model ID for the base model (required)
- `--output-dir`: Directory to save model (required)
- `--api-key`: HuggingFace API key (optional, will use SIA_HF_API_KEY)
### train_mistral
Fine-tune Mistral models using Mistral's API.
```bash
train_mistral --model mistral-large-latest
```
Options:
- `--config`: Path to training configuration file (default: /root/sia/training/config.yaml)
- `--model`: Base model name (default: mistral-large-latest)
- `--api-key`: Mistral API key (optional, will use SIA_MISTRAL_API_KEY)
## Configuration Format
The training configuration file (YAML) should include:
```yaml
model:
system_prompt_path: "/root/sia/system_prompt.md"
action_schema: "/root/sia/action_schema.xsd"
params:
learning_rate: 1e-5
epochs: 3
data:
- "/root/sia/training/data_dir1/"
- "/root/sia/training/data_dir2/"
```
## Data Format
Training data should be XML files in the following format:
```xml
<iteration system_prompt_hash="..." action_schema_hash="...">
<context>
<!-- XML context -->
</context>
<response>
<!-- Model response -->
</response>
</iteration>
```

13
tools/train/requirements.txt Normal file
View File

@@ -0,0 +1,13 @@
pyyaml>=6.0
requests>=2.28.0
torch>=2.0.0
transformers>=4.30.0
# DeepSeek support
accelerate>=0.25.0
bitsandbytes>=0.41.1
einops>=0.7.0
sentencepiece>=0.1.99
unsloth>=2024.3
trl>=0.7.8
datasets>=2.14.6
peft>=0.8.0

36
tools/train/setup.py Normal file
View File

@@ -0,0 +1,36 @@
from setuptools import setup, find_packages
setup(
name="train",
version="0.1.0",
packages=find_packages(),
scripts=[
'bin/train_deepseek',
'bin/train_mistral'
],
install_requires=[
'pyyaml>=6.0',
'requests>=2.28.0',
'torch>=2.0.0',
'transformers>=4.30.0',
'accelerate>=0.25.0',
'bitsandbytes>=0.41.1',
'einops>=0.7.0',
'sentencepiece>=0.1.99',
'unsloth>=2024.3',
'trl>=0.7.8',
'datasets>=2.14.6',
'peft>=0.8.0',
'pytest>=7.0.0',
'pytest-cov>=4.0.0',
'black>=22.0.0',
'flake8>=4.0.0'
],
classifiers=[
'Development Status :: 3 - Alpha',
'Intended Audience :: Developers',
'Programming Language :: Python :: 3',
'Programming Language :: Python :: 3.10',
],
python_requires='>=3.10',
)

15
tools/train/train.sh Normal file
View File

@@ -0,0 +1,15 @@
#!/bin/bash
set -e
SIA_DIR="/root/sia"
OUTPUT_DIR="${1:-/root/models/$(cd "$SIA_DIR" && git rev-parse HEAD)}"
if [ -n "$(cd "$SIA_DIR" && git status --porcelain)" ]; then
echo "Uncommitted changes in SIA directory"
#exit 1
fi
mkdir -p "$OUTPUT_DIR"
train_deepseek --output-dir "$OUTPUT_DIR"

8
tools/train/train/__init__.py Normal file
View File

@@ -0,0 +1,8 @@
"""
SIA Training Tool
This package provides utilities for fine-tuning language models used by SIA.
Supports DeepSeek and Mistral models.
"""
__version__ = "0.1.0"

141
tools/train/train/mistral_api.py Normal file
View File

@@ -0,0 +1,141 @@
#!/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())

239
tools/train/train/unsloth_deepseek.py Normal file
View File

@@ -0,0 +1,239 @@
#!/root/venvs/train/bin/python
"""
Script for fine-tuning DeepSeek models for SIA using Unsloth.
Training always starts from a base model and creates a new fine-tuned model.
"""
import argparse
import os
import sys
import torch
from dataclasses import dataclass
from pathlib import Path
import json
# Import from shared library
from .util import prepare_training_data
@dataclass
class Config:
def __init__(self):
parser = argparse.ArgumentParser(description='Train SIA model using Unsloth')
parser.add_argument(
'--config',
type=Path,
default=Path('/root/sia/training/config.yaml'),
help='Path to config file'
)
parser.add_argument(
'--base-model',
type=str,
default='deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B',
help='HuggingFace model ID for base model'
)
parser.add_argument(
'--output-dir',
type=Path,
required=True,
help='Directory to save the trained model'
)
parser.add_argument(
'--api-key',
type=str,
default=os.environ.get('SIA_HF_API_KEY'),
help='HuggingFace API key'
)
self.args = parser.parse_args()
@property
def config_path(self) -> Path:
return self.args.config
@property
def base_model(self) -> str:
return self.args.base_model
@property
def output_dir(self) -> Path:
return self.args.output_dir
@property
def api_key(self) -> str:
return self.args.api_key
def train_model(config: Config, training_data, train_params, commit_hash):
"""Train the model using Unsloth"""
try:
from unsloth import FastLanguageModel
from transformers import TrainingArguments, DataCollatorForSeq2Seq
from trl import SFTTrainer
from datasets import Dataset
from unsloth.chat_templates import get_chat_template, train_on_responses_only
except ImportError as e:
print(f"Error importing required libraries: {e}")
print("Please ensure Unsloth and its dependencies are installed.")
sys.exit(1)
print(f"Starting training from base model: {config.base_model}")
# Convert to datasets format
dataset = Dataset.from_list(training_data)
# Determine if bfloat16 is supported
dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
# Load the model - always from a base model (no incremental updates)
try:
model, tokenizer = FastLanguageModel.from_pretrained(
model_name=config.base_model,
max_seq_length=2048,
dtype=dtype,
load_in_4bit=True,
token=config.api_key,
)
except Exception as e:
print(f"Error loading base model: {e}")
sys.exit(1)
# Apply LoRA
target_modules = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
model = FastLanguageModel.get_peft_model(
model,
r=16,
target_modules=target_modules,
lora_alpha=16,
lora_dropout=0,
bias="none",
use_gradient_checkpointing="unsloth",
random_state=3407,
)
# Apply chat template
tokenizer = get_chat_template(
tokenizer,
chat_template="llama-3.1", # Compatible with DeepSeek
)
# Function to format conversations
def formatting_prompts_func(examples):
convos = examples["conversations"]
texts = [tokenizer.apply_chat_template(convo, tokenize=False, add_generation_prompt=False) for convo in convos]
return {"text": texts}
# Standarize dataset and format
from unsloth.chat_templates import standardize_sharegpt
# Add conversations field if not present
if "conversations" not in dataset.column_names:
if "messages" in dataset.column_names:
dataset = dataset.rename_column("messages", "conversations")
else:
dataset = dataset.map(lambda x: {"conversations": [{"role": "system", "content": x.get("system_prompt", "")},
{"role": "user", "content": x.get("prompt", "")},
{"role": "assistant", "content": x.get("response", "")}]})
# Standardize format
dataset = standardize_sharegpt(dataset)
# Apply formatting
dataset = dataset.map(formatting_prompts_func, batched=True)
# Configure the trainer
output_dir = config.output_dir / commit_hash
output_dir.mkdir(parents=True, exist_ok=True)
# Determine steps or epochs based on dataset size
max_steps = None
num_train_epochs = train_params.epochs
if len(dataset) < 100: # Small dataset
# Aim for at least 500 steps for small datasets
max_steps = 500
num_train_epochs = None
trainer = SFTTrainer(
model=model,
tokenizer=tokenizer,
train_dataset=dataset,
dataset_text_field="text",
max_seq_length=2048,
data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer),
dataset_num_proc=2,
packing=False,
args=TrainingArguments(
per_device_train_batch_size=2,
gradient_accumulation_steps=4,
warmup_steps=5,
max_steps=max_steps,
num_train_epochs=num_train_epochs,
learning_rate=train_params.learning_rate,
fp16=not torch.cuda.is_bf16_supported(),
bf16=torch.cuda.is_bf16_supported(),
logging_steps=10,
optim="adamw_8bit",
weight_decay=0.01,
lr_scheduler_type="linear",
seed=3407,
output_dir=str(output_dir),
report_to="none",
),
)
# Train only on responses
trainer = train_on_responses_only(
trainer,
instruction_part="<|start_header_id|>user<|end_header_id|>\n\n",
response_part="<|start_header_id|>assistant<|end_header_id|>\n\n",
)
# Train the model
trainer.train()
# Enable inference mode for the model
model = FastLanguageModel.for_inference(model)
# Save the model
model.save_pretrained(output_dir)
tokenizer.save_pretrained(output_dir)
# Create a metadata file with training information
with open(output_dir / "training_info.json", "w") as f:
json.dump({
"base_model": config.base_model,
"commit_hash": commit_hash,
"learning_rate": train_params.learning_rate,
"epochs": train_params.epochs,
"dataset_size": len(dataset),
"training_method": "unsloth",
}, f, indent=2)
return output_dir
def main():
config = Config()
# Prepare training data
training_data, train_params, commit_hash = prepare_training_data(config.config_path)
if not training_data:
print("No valid training data found. Exiting.")
return 1
# Train the model
try:
model_dir = train_model(config, training_data, train_params, commit_hash)
# Create symlink to current
current_link = config.output_dir / "current"
if current_link.exists() or current_link.is_symlink():
current_link.unlink()
os.symlink(model_dir, current_link, target_is_directory=True)
print(f"Training complete. Model saved to {model_dir}")
print(f"Symlink created at {current_link}")
return 0
except Exception as e:
print(f"Error during training: {e}")
return 1
if __name__ == "__main__":
exit(main())

186
tools/train/train/util.py Normal file
View File

@@ -0,0 +1,186 @@
"""
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
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 = 1
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, str]:
"""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)
paths = list(xml_files)
paths.append(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)
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(
learning_rate=config_data['params'].get('learning_rate', 1e-5),
epochs=config_data['params'].get('epochs', 3),
batch_size=config_data['params'].get('batch_size', 1)
)
return training_data, train_params, commit_hash
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}")

260
tools/train/train_mistral.py
View File

@@ -1,260 +0,0 @@
#!/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())