Create dataset_gen.py

dataset_gen.py

@@ -0,0 +1,37 @@
+import json
+import os
+from parser import parse_source_to_graph
+from datetime import datetime
+
+def create_hf_dataset(code_samples, output_file="software_structure_dataset.jsonl"):
+    """
+    Takes a list of code strings, parses them, and saves them
+    in a format ready for Hugging Face 'datasets'.
+    """
+    data_entries = []
+    
+    for idx, code in enumerate(code_samples):
+        graph_data = parse_source_to_graph(code)
+        
+        if "error" in graph_data:
+            continue
+            
+        # Flatten vectors for ML input
+        vectors = [n['vector'] for n in graph_data['nodes']]
+        
+        entry = {
+            "id": f"sample_{idx}_{int(datetime.now().timestamp())}",
+            "source_code": code,
+            "graph_nodes": json.dumps(graph_data['nodes']),
+            "graph_edges": json.dumps(graph_data['connections']),
+            "structural_vectors": vectors, # The core feature for training
+            "node_count": len(graph_data['nodes'])
+        }
+        data_entries.append(entry)
+
+    # Write to JSONL
+    with open(output_file, 'w') as f:
+        for entry in data_entries:
+            f.write(json.dumps(entry) + '\n')
+            
+    return output_file