134 lines
4.8 KiB
Python
134 lines
4.8 KiB
Python
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.glob('*.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() |