Update dataset_gen.py

dataset_gen.py

@@ -2,37 +2,74 @@ import json
 import os
 from parser import parse_source_to_graph
 from datetime import datetime
+from huggingface_hub import HfApi
 
 OUTPUT_FILE = "pystructure_dataset.jsonl"
 
 def create_dataset_entry(code):
-    """
-    Parses code and appends a training example to the JSONL file.
-    """
     graph_data = parse_source_to_graph(code)
     
     if "error" in graph_data:
         return {"status": "error", "message": graph_data["error"]}
         
-    vectors = [n['vector'] for n in graph_data['nodes']]
+    vectors = [n['vec'] 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,
+        "source_code": code, # We keep full source for training
         "meta": {
             "node_count": len(graph_data['nodes']),
-            "max_depth": max([n['level'] for n in graph_data['nodes']]) if graph_data['nodes'] else 0
+            "max_depth": max([n['lvl'] for n in graph_data['nodes']]) if graph_data['nodes'] else 0,
+            "snippet": code[:50].replace('\n', ' ') + "..." # For UI preview
+        },
+        # Store compact structure for training
+        "structure": {
+            "vectors": vectors,
+            "edges": graph_data['connections']
         }
     }
 
-    # Append to JSONL file
     with open(OUTPUT_FILE, 'a') as f:
         f.write(json.dumps(entry) + '\n')
             
-    return {"status": "success", "file": OUTPUT_FILE, "entry_id": entry['id']}
+    return {"status": "success", "id": entry['id']}
+
+def get_dataset_stats():
+    """Reads metadata from the JSONL file without loading heavy source code."""
+    entries = []
+    if not os.path.exists(OUTPUT_FILE):
+        return []
+        
+    with open(OUTPUT_FILE, 'r') as f:
+        for line in f:
+            try:
+                data = json.loads(line)
+                # Only return lightweight info for the UI table
+                entries.append({
+                    "id": data['id'],
+                    "timestamp": data['timestamp'],
+                    "node_count": data['meta']['node_count'],
+                    "snippet": data['meta']['snippet']
+                })
+            except:
+                continue
+    return entries[::-1] # Newest first
+
+def upload_to_hub(token, repo_id):
+    """Pushes the local JSONL file to Hugging Face."""
+    if not os.path.exists(OUTPUT_FILE):
+        return {"status": "error", "message": "No dataset found."}
+        
+    try:
+        api = HfApi(token=token)
+        # Upload the specific file
+        api.upload_file(
+            path_or_fileobj=OUTPUT_FILE,
+            path_in_repo="dataset.jsonl",
+            repo_id=repo_id,
+            repo_type="dataset"
+        )
+        return {"status": "success", "message": f"Uploaded to https://huggingface.co/datasets/{repo_id}"}
+    except Exception as e:
+        return {"status": "error", "message": str(e)}