Update dataset_gen.py

dataset_gen.py

@@ -7,23 +7,26 @@ from huggingface_hub import HfApi
 OUTPUT_FILE = "pystructure_dataset.jsonl"
 
 def create_dataset_entry(code):
+    """
+    Parses code and appends a training example to the local JSONL file.
+    """
     graph_data = parse_source_to_graph(code)
     
     if "error" in graph_data:
         return {"status": "error", "message": graph_data["error"]}
         
+    # Flatten vectors for ML input
     vectors = [n['vec'] for n in graph_data['nodes']]
     
     entry = {
         "id": f"sample_{int(datetime.now().timestamp())}",
         "timestamp": datetime.now().isoformat(),
-        "source_code": code, # We keep full source for training
+        "source_code": code,
         "meta": {
             "node_count": len(graph_data['nodes']),
             "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
+            "snippet": code[:50].replace('\n', ' ') + "..."
         },
-        # Store compact structure for training
         "structure": {
             "vectors": vectors,
             "edges": graph_data['connections']
@@ -36,7 +39,9 @@ def create_dataset_entry(code):
     return {"status": "success", "id": entry['id']}
 
 def get_dataset_stats():
-    """Reads metadata from the JSONL file without loading heavy source code."""
+    """
+    Reads metadata from the local JSONL file for the UI table.
+    """
     entries = []
     if not os.path.exists(OUTPUT_FILE):
         return []
@@ -45,7 +50,6 @@ def get_dataset_stats():
         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'],
@@ -54,22 +58,64 @@ def get_dataset_stats():
                 })
             except:
                 continue
-    return entries[::-1] # Newest first
+    return entries[::-1] # Return newest first
 
-def upload_to_hub(token, repo_id):
-    """Pushes the local JSONL file to Hugging Face."""
+def upload_to_hub(token, repo_name_input):
+    """
+    1. Autodetects username from token.
+    2. Creates repo if it doesn't exist.
+    3. Uploads the local file as a unique shard to 'append' to the dataset.
+    """
     if not os.path.exists(OUTPUT_FILE):
-        return {"status": "error", "message": "No dataset found."}
+        return {"status": "error", "message": "No local dataset found to upload."}
         
     try:
         api = HfApi(token=token)
-        # Upload the specific file
+
+        # 1. Auto-detect Username
+        try:
+            user_info = api.whoami()
+            username = user_info['name']
+        except Exception:
+            return {"status": "error", "message": "Invalid HF Token. Please check your write token."}
+
+        # 2. Resolve Repo ID
+        # If user typed "my-dataset", convert to "username/my-dataset"
+        # If user typed "username/my-dataset", keep it as is
+        if "/" in repo_name_input:
+            full_repo_id = repo_name_input
+        else:
+            full_repo_id = f"{username}/{repo_name_input}"
+
+        # 3. Create Repo (Idempotent)
+        # exist_ok=True prevents errors if the repo already exists
+        api.create_repo(
+            repo_id=full_repo_id, 
+            repo_type="dataset", 
+            exist_ok=True
+        )
+
+        # 4. Upload with Sharding (Simulated Append)
+        # We assume the local file contains *new* data. 
+        # We upload it with a unique timestamp filename.
+        # HF 'datasets' library will automatically load ALL jsonl files in the directory.
+        timestamp = int(datetime.now().timestamp())
+        remote_filename = f"data_shard_{timestamp}.jsonl"
+
         api.upload_file(
             path_or_fileobj=OUTPUT_FILE,
-            path_in_repo="dataset.jsonl",
-            repo_id=repo_id,
+            path_in_repo=remote_filename,
+            repo_id=full_repo_id,
             repo_type="dataset"
         )
-        return {"status": "success", "message": f"Uploaded to https://huggingface.co/datasets/{repo_id}"}
+        
+        # Optional: Rename local file to avoid re-uploading duplicate data next time?
+        # For now, we leave it as is, but this logic implies the user manages their local file.
+        
+        return {
+            "status": "success", 
+            "message": f"Successfully appended data to https://huggingface.co/datasets/{full_repo_id} (File: {remote_filename})"
+        }
+
     except Exception as e:
         return {"status": "error", "message": str(e)}