Update run_job.py

run_job.py

@@ -1,188 +1,204 @@
-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_2.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",
-            "--phrase-lexicon", "sys7_phrase_lexicons_desc_only.json",
-            "--auto-language-detect",
-            "--language-allowlist", "en",
-            "--min-chars", "10",
-            "--min-tokens", "2",
-            "--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)
+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]
+        parquet_files = sorted(parquet_files)
+        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"
+
+    # Resume support: skip shards already uploaded to OUTPUT_REPO.
+    existing_outputs = set()
+    try:
+        existing_outputs = set(api.list_repo_files(repo_id=OUTPUT_REPO, repo_type="dataset"))
+        done = [p for p in existing_outputs if p.startswith("data/sys7_features_part_") and p.endswith(".parquet")]
+        log(f"Resume: found {len(done)} existing output shards in {OUTPUT_REPO}.")
+    except Exception as e:
+        log(f"Resume warning: could not list output repo files: {e}")
+    
+    for i, file_path in enumerate(parquet_files):
+        shard_output_name = f"sys7_features_part_{i:04d}.parquet"
+        remote_out_path = f"data/{shard_output_name}"
+        if remote_out_path in existing_outputs:
+            log(f"Skipping Shard {i+1}/{len(parquet_files)} (already uploaded): {remote_out_path}")
+            continue
+
+        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_path = os.path.join(output_dir, shard_output_name)
+
+        cmd = [
+            "python", "-u", "sys7_miner_2.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",
+            "--phrase-lexicon", "sys7_phrase_lexicons_desc_only.json",
+            "--auto-language-detect",
+            "--language-allowlist", "en",
+            "--min-chars", "10",
+            "--min-tokens", "2",
+            "--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"
+            )
+            existing_outputs.add(remote_out_path)
+            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)