Upload run_job.py

run_job.py

@@ -0,0 +1,183 @@
+import gradio as gr
+import os
+import threading
+import subprocess
+import time
+import shutil
+from huggingface_hub import HfApi, hf_hub_download, create_repo
+
+# --- CONFIG ---
+INPUT_DATASET = "The-data-company/TikTok-10M"
+# CHANGE THIS TO YOUR USERNAME!
+OUTPUT_REPO = "Imaginethat/tiktok-sys7-mined-results" 
+HF_TOKEN = os.environ.get("HF_TOKEN")
+
+# Global variables for the UI
+log_buffer = "Initializing Environment...\n"
+dataset_link_md = ""
+mining_active = False
+
+def log(msg):
+    """Adds a message to the UI log window"""
+    global log_buffer
+    timestamp = time.strftime("%H:%M:%S")
+    entry = f"[{timestamp}] {msg}"
+    print(entry) # Keep console logs as backup
+    log_buffer += entry + "\n"
+
+def run_mining_logic():
+    global dataset_link_md, mining_active
+    
+    if not HF_TOKEN:
+        log("ERROR: HF_TOKEN secret is missing in Settings!")
+        return
+
+    api = HfApi(token=HF_TOKEN)
+    
+    # 1. Setup Repo
+    log(f"Setting up Output Repo: {OUTPUT_REPO}")
+    try:
+        create_repo(repo_id=OUTPUT_REPO, repo_type="dataset", exist_ok=True)
+    except Exception as e:
+        log(f"Repo setup note: {e}")
+
+    # 2. List Files
+    log(f"Scanning {INPUT_DATASET} for files...")
+    try:
+        all_files = api.list_repo_files(repo_id=INPUT_DATASET, repo_type="dataset")
+        parquet_files = [f for f in all_files if f.endswith(".parquet") and "train" in f]
+        log(f"Found {len(parquet_files)} shards to process.")
+    except Exception as e:
+        log(f"Failed to list files: {e}")
+        return
+
+    # 3. Process Shards
+    input_dir = "./data_input"
+    output_dir = "./data_output"
+    
+    for i, file_path in enumerate(parquet_files):
+        log(f"--- Processing Shard {i+1}/{len(parquet_files)}: {file_path} ---")
+        
+        # Cleanup
+        if os.path.exists(input_dir): shutil.rmtree(input_dir)
+        if os.path.exists(output_dir): shutil.rmtree(output_dir)
+        os.makedirs(input_dir, exist_ok=True)
+        os.makedirs(output_dir, exist_ok=True)
+
+        # Download
+        try:
+            local_path = hf_hub_download(
+                repo_id=INPUT_DATASET,
+                filename=file_path,
+                repo_type="dataset",
+                local_dir=input_dir,
+                token=HF_TOKEN
+            )
+        except Exception as e:
+            log(f"Download failed: {e}")
+            continue
+
+        # Run Miner (Streaming output to UI)
+        shard_output_name = f"sys7_features_part_{i:04d}.parquet"
+        shard_output_path = os.path.join(output_dir, shard_output_name)
+
+        cmd = [
+            "python", "-u", "sys7_miner.py", # -u forces unbuffered output
+            "--input-path", input_dir,
+            "--output-parquet", shard_output_path,
+            "--lexicons", "system7_lexicons.json",
+            "--label-orders", "label_orders.json",
+            "--slang-lexicon", "slang_lexicon.json",
+            "--device", "cuda",
+            "--batch-size", "50000",
+            "--embedding-batch-size", "512",
+            "--workers", "4"
+        ]
+
+        # Use Popen to capture logs in real-time
+        try:
+            process = subprocess.Popen(
+                cmd, 
+                stdout=subprocess.PIPE, 
+                stderr=subprocess.STDOUT, 
+                text=True, 
+                bufsize=1
+            )
+            
+            # Read stdout line by line
+            for line in process.stdout:
+                if "Processed" in line or "Error" in line: # Filter logs to keep UI clean?
+                    log("MINER: " + line.strip())
+                elif "Loading" in line:
+                    log("MINER: " + line.strip())
+            
+            process.wait()
+            
+            if process.returncode != 0:
+                log(f"Miner failed with code {process.returncode}")
+                continue
+                
+        except Exception as e:
+            log(f"Subprocess error: {e}")
+            continue
+
+        # Upload
+        try:
+            log(f"Uploading {shard_output_name}...")
+            api.upload_file(
+                path_or_fileobj=shard_output_path,
+                path_in_repo=f"data/{shard_output_name}",
+                repo_id=OUTPUT_REPO,
+                repo_type="dataset"
+            )
+            log("Upload Successful!")
+        except Exception as e:
+            log(f"Upload failed: {e}")
+
+    # 4. Finish
+    final_url = f"https://huggingface.co/datasets/{OUTPUT_REPO}"
+    log("ALL JOBS COMPLETE.")
+    dataset_link_md = f"## 🎉 DONE! Data is ready: [Click Here to View Dataset]({final_url})"
+    mining_active = False
+
+def start_thread():
+    global mining_active
+    if mining_active:
+        return "Already Running..."
+    mining_active = True
+    t = threading.Thread(target=run_mining_logic)
+    t.start()
+    return "Mining Started! Watch logs below."
+
+# --- GRADIO UI ---
+with gr.Blocks(title="TikTok Miner") as demo:
+    gr.Markdown("# ⛏️ System 7 Miner Control")
+    
+    with gr.Row():
+        start_btn = gr.Button("🚀 Start Mining Job", variant="primary")
+        status_txt = gr.Textbox(label="Status", value="Ready to start.", interactive=False)
+
+    # The Log Window
+    logs_out = gr.TextArea(
+        label="Live Process Logs", 
+        placeholder="Logs will stream here...", 
+        lines=20,
+        max_lines=25,
+        autoscroll=True
+    )
+    
+    # The Success Link (Hidden until done)
+    link_display = gr.Markdown("")
+
+    # Button Action
+    start_btn.click(fn=start_thread, inputs=None, outputs=status_txt)
+
+    # Auto-refresh the logs every 1 second
+    def update_ui():
+        return log_buffer, dataset_link_md
+        
+    timer = gr.Timer(1)
+    timer.tick(fn=update_ui, outputs=[logs_out, link_display])
+
+# Launch on the port HF expects
+demo.launch(server_name="0.0.0.0", server_port=7860)