diff --git a/finetune_dataset_creator.py b/finetune_dataset_creator.py new file mode 100644 index 0000000..3c1ba1e --- /dev/null +++ b/finetune_dataset_creator.py @@ -0,0 +1,133 @@ +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 or + root.get('action_schema_hash') != self.action_schema_hash): + print(f"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() \ No newline at end of file