Update dataset_gen.py

dataset_gen.py

@@ -3,35 +3,36 @@ 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"):
+OUTPUT_FILE = "pystructure_dataset.jsonl"
+
+def create_dataset_entry(code):
     """
-    Takes a list of code strings, parses them, and saves them
-    in a format ready for Hugging Face 'datasets'.
+    Parses code and appends a training example to the JSONL file.
     """
-    data_entries = []
+    graph_data = parse_source_to_graph(code)
     
-    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']]
+    if "error" in graph_data:
+        return {"status": "error", "message": graph_data["error"]}
         
-        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'])
+    vectors = [n['vector'] for n in graph_data['nodes']]
+    
+    entry = {
+        "id": f"sample_{int(datetime.now().timestamp())}",
+        "timestamp": datetime.now().isoformat(),
+        "source_code": code,
+        "graph_structure": {
+            "nodes": [n['id'] for n in graph_data['nodes']],
+            "edges": graph_data['connections']
+        },
+        "structural_vectors": vectors,
+        "meta": {
+            "node_count": len(graph_data['nodes']),
+            "max_depth": max([n['level'] for n in graph_data['nodes']]) if graph_data['nodes'] else 0
         }
-        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')
+    # Append to JSONL file
+    with open(OUTPUT_FILE, 'a') as f:
+        f.write(json.dumps(entry) + '\n')
             
-    return output_file
+    return {"status": "success", "file": OUTPUT_FILE, "entry_id": entry['id']}