v2

app.py

@@ -2,10 +2,18 @@ import os
 import gradio as gr
 import duckdb
 import requests
+import threading
+import time
+import json
+from datetime import datetime
 
 DATASET_REPO = "gmongaras/Imagenet21K"
 HF_TOKEN = os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_HUB_TOKEN")
 
+# 全局状态文件
+STATUS_FILE = "extraction_status.json"
+OUTPUT_FILE = "ids.parquet"
+
 def get_parquet_urls():
     """获取所有 parquet 文件 URL"""
     api_url = f"https://huggingface.co/api/datasets/{DATASET_REPO}/parquet/default/train"
@@ -28,63 +36,137 @@ def get_parquet_urls():
     
     return urls
 
-def build_ids_duckdb(progress=gr.Progress()):
-    """使用 DuckDB 直接查询所有 Parquet 文件"""
-    
-    progress(0, desc="获取 Parquet 文件列表...")
-    urls = get_parquet_urls()
+def save_status(status, message, progress=0.0, error=None):
+    """保存状态到文件"""
+    status_data = {
+        "status": status,  # "running", "completed", "error", "idle"
+        "message": message,
+        "progress": progress,
+        "timestamp": datetime.now().isoformat(),
+        "error": error
+    }
+    
+    with open(STATUS_FILE, 'w') as f:
+        json.dump(status_data, f)
+
+def load_status():
+    """从文件加载状态"""
+    if os.path.exists(STATUS_FILE):
+        try:
+            with open(STATUS_FILE, 'r') as f:
+                return json.load(f)
+        except:
+            pass
+    
+    return {
+        "status": "idle",
+        "message": "没有运行中的任务",
+        "progress": 0.0,
+        "timestamp": "",
+        "error": None
+    }
+
+def build_ids_background():
+    """后台运行 DuckDB 提取任务"""
+    try:
+        save_status("running", "获取 Parquet 文件列表...", 0.0)
+        
+        urls = get_parquet_urls()
+        save_status("running", f"找到 {len(urls)} 个文件，开始提取...", 0.1)
+        
+        con = duckdb.connect()
+        con.execute("INSTALL httpfs; LOAD httpfs;")
+        
+        # 设置认证和优化
+        if HF_TOKEN:
+            con.execute(f"SET httpfs_custom_header='Authorization: Bearer {HF_TOKEN}';")
+        
+        con.execute("SET httpfs_keep_alive=true;")
+        con.execute("SET enable_object_cache=true;")
+        con.execute("SET threads=4;")
+        
+        save_status("running", "执行 SQL 查询...", 0.2)
+        
+        # 构建文件列表
+        files_literal = ",".join([f"'{url}'" for url in urls])
+        
+        # 一次性查询所有文件的 id 列
+        query = f"""
+        COPY (
+            SELECT id 
+            FROM parquet_scan([{files_literal}])
+        ) TO '{OUTPUT_FILE}' (FORMAT 'parquet', COMPRESSION 'zstd');
+        """
+        
+        save_status("running", "正在执行大规模查询...", 0.5)
+        con.execute(query)
+        
+        # 完成
+        file_size = os.path.getsize(OUTPUT_FILE) / 1024 / 1024
+        save_status("completed", f"提取完成！文件大小: {file_size:.1f} MB", 1.0)
+        
+    except Exception as e:
+        save_status("error", "提取失败", 0.0, str(e))
+
+def start_extraction():
+    """启动后台提取任务"""
+    status = load_status()
     
-    progress(0.1, desc=f"找到 {len(urls)} 个文件，开始提取...")
+    if status["status"] == "running":
+        return "任务已在运行中..."
     
-    con = duckdb.connect()
-    con.execute("INSTALL httpfs; LOAD httpfs;")
+    # 删除旧的输出文件
+    if os.path.exists(OUTPUT_FILE):
+        os.remove(OUTPUT_FILE)
     
-    # 设置认证
-    if HF_TOKEN:
-        con.execute(f"SET httpfs_custom_header='Authorization: Bearer {HF_TOKEN}';")
+    # 启动后台线程
+    thread = threading.Thread(target=build_ids_background, daemon=True)
+    thread.start()
     
-    progress(0.2, desc="执行 SQL 查询...")
+    return "后台任务已启动！刷新页面不会中断任务。"
+
+def check_status():
+    """检查当前状态"""
+    status = load_status()
     
-    # 构建文件列表
-    files_literal = ",".join([f"'{url}'" for url in urls])
+    progress_text = f"{status['progress']*100:.1f}%" if status['progress'] > 0 else ""
+    timestamp = status.get('timestamp', '')
     
-    # 一次性查询所有文件的 id 列
-    query = f"""
-    COPY (
-        SELECT id 
-        FROM parquet_scan([{files_literal}])
-    ) TO 'ids.parquet' (FORMAT 'parquet', COMPRESSION 'zstd');
-    """
+    if status['status'] == 'completed':
+        if os.path.exists(OUTPUT_FILE):
+            return f"✅ {status['message']}\n时间: {timestamp}\n进度: {progress_text}", OUTPUT_FILE
+        else:
+            return f"⚠️ 任务完成但文件不存在\n时间: {timestamp}", None
     
-    progress(0.5, desc="正在执行大规模查询...")
-    con.execute(query)
+    elif status['status'] == 'error':
+        error_msg = status.get('error', '未知错误')
+        return f"❌ {status['message']}\n错误: {error_msg}\n时间: {timestamp}", None
     
-    progress(1.0, desc="完成！")
+    elif status['status'] == 'running':
+        return f"🔄 {status['message']}\n进度: {progress_text}\n时间: {timestamp}", None
     
-    file_size = os.path.getsize('ids.parquet') / 1024 / 1024
-    return f"提取完成！\n文件大小: {file_size:.1f} MB"
-
-def ui_build_ids(progress=gr.Progress()):
-    try:
-        result = build_ids_duckdb(progress=progress)
-        return result, "ids.parquet"
-    except Exception as e:
-        return f"错误: {e}", None
+    else:
+        return status['message'], None
 
 with gr.Blocks() as demo:
-    gr.Markdown("# ImageNet21K ID 提取器 (DuckDB)\n使用 DuckDB 快速提取 ID 列")
+    gr.Markdown("# ImageNet21K ID 提取器 (后台运行)\n✨ 支持浏览器断线重连，任务在后台持续运行")
     
-    run_btn = gr.Button("开始提取", variant="primary")
+    with gr.Row():
+        start_btn = gr.Button("开始提取", variant="primary")
+        refresh_btn = gr.Button("刷新状态", variant="secondary")
     
     with gr.Row():
-        log = gr.Textbox(label="状态", lines=5)
-        download = gr.File(label="下载文件")
-    
-    run_btn.click(
-        ui_build_ids,
-        outputs=[log, download],
-        show_progress=True
-    )
+        status_text = gr.Textbox(label="任务状态", lines=5)
+        download_file = gr.File(label="下载文件")
+    
+    # 自动刷新状态 (每10秒)
+    auto_refresh = gr.Textbox(visible=False)
+    
+    start_btn.click(start_extraction, outputs=[status_text])
+    refresh_btn.click(check_status, outputs=[status_text, download_file])
+    
+    # 页面加载时检查状态
+    demo.load(check_status, outputs=[status_text, download_file])
 
 demo.queue()