Add application file

app.py

@@ -0,0 +1,92 @@
+import os
+import gradio as gr
+import duckdb
+import requests
+
+DATASET_REPO = "gmongaras/Imagenet21K"
+HF_TOKEN = os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_HUB_TOKEN")
+
+def get_parquet_urls():
+    """获取所有 parquet 文件 URL"""
+    api_url = f"https://huggingface.co/api/datasets/{DATASET_REPO}/parquet/default/train"
+    headers = {"Authorization": f"Bearer {HF_TOKEN}"} if HF_TOKEN else {}
+    
+    response = requests.get(api_url, headers=headers)
+    response.raise_for_status()
+    
+    data = response.json()
+    urls = []
+    
+    items = data if isinstance(data, list) else data.get("parquet_files", [])
+    for item in items:
+        if isinstance(item, str):
+            urls.append(item)
+        else:
+            url = item.get("url")
+            if url:
+                urls.append(url)
+    
+    return urls
+
+def build_ids_duckdb(progress=gr.Progress()):
+    """使用 DuckDB 直接查询所有 Parquet 文件"""
+    
+    progress(0, desc="获取 Parquet 文件列表...")
+    urls = get_parquet_urls()
+    
+    progress(0.1, desc=f"找到 {len(urls)} 个文件，开始提取...")
+    
+    con = duckdb.connect()
+    con.execute("INSTALL httpfs; LOAD httpfs;")
+    
+    # 设置认证
+    if HF_TOKEN:
+        con.execute(f"SET httpfs_custom_header='Authorization: Bearer {HF_TOKEN}';")
+    
+    progress(0.2, desc="执行 SQL 查询...")
+    
+    # 构建文件列表
+    files_literal = ",".join([f"'{url}'" for url in urls])
+    
+    # 一次性查询所有文件的 id 列
+    query = f"""
+    COPY (
+        SELECT id 
+        FROM parquet_scan([{files_literal}])
+    ) TO 'ids.parquet' (FORMAT 'parquet', COMPRESSION 'zstd');
+    """
+    
+    progress(0.5, desc="正在执行大规模查询...")
+    con.execute(query)
+    
+    progress(1.0, desc="完成！")
+    
+    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
+
+with gr.Blocks() as demo:
+    gr.Markdown("# ImageNet21K ID 提取器 (DuckDB)\n使用 DuckDB 快速提取 ID 列")
+    
+    run_btn = gr.Button("开始提取", variant="primary")
+    
+    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
+    )
+
+demo.queue()
+
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -0,0 +1,3 @@
+gradio
+requests
+duckdb