Added dataset creator

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()