#!/usr/bin/env python3 """MinerU Gradio 客户端 — 纯前端,所有解析任务通过 HTTP 转发到 Router/Worker。 与 mineru-gradio 的区别: - 不启动内嵌 FastAPI server - 不需要 GPU / vLLM - 通过 HTTP 调用 mineru-router → mineru-worker 处理任务 """ import os import time import asyncio import tempfile import zipfile import json from pathlib import Path import gradio as gr import httpx API_BASE = os.environ.get("MINERU_API_BASE", "http://mineru-router:8000") DEFAULT_BACKEND = os.environ.get("MINERU_BACKEND", "hybrid-auto-engine") DEFAULT_LANG = os.environ.get("MINERU_LANG", "ch") POLL_INTERVAL = float(os.environ.get("MINERU_POLL_INTERVAL", "1.0")) MAX_WAIT = float(os.environ.get("MINERU_MAX_WAIT", "600")) async def discover_api(client: httpx.AsyncClient) -> str: """探测实际可用的 API 前缀: /api/v1/tasks 或 /tasks""" candidates = [ f"{API_BASE}/api/v1/health", f"{API_BASE}/health", f"{API_BASE}/docs", ] for url in candidates: try: r = await client.get(url, timeout=5) if r.status_code < 500: # FastAPI docs → /api/v1 ; plain health → / if "/api/v1" in str(r.url): return "/api/v1/tasks" return "/tasks" except Exception: continue # 默认新版 API return "/api/v1/tasks" async def submit_task(client: httpx.AsyncClient, api_prefix: str, file_path: str, file_name: str, backend: str, lang: str) -> dict: """提交解析任务""" submit_url = f"{API_BASE}{api_prefix}/submit" with open(file_path, "rb") as f: files = {"file": (file_name, f, "application/pdf")} data = { "backend": backend, "lang": lang, } r = await client.post(submit_url, files=files, data=data, timeout=30) r.raise_for_status() return r.json() async def get_task_status(client: httpx.AsyncClient, api_prefix: str, task_id: str) -> dict: """查询任务状态""" url = f"{API_BASE}{api_prefix}/{task_id}" r = await client.get(url, timeout=10) r.raise_for_status() return r.json() async def download_result(client: httpx.AsyncClient, api_prefix: str, task_id: str, output_dir: str) -> str: """下载任务结果。尝试多种 endpoint 格式""" candidates = [ f"{api_prefix}/{task_id}/data", f"{api_prefix}/{task_id}/result", ] for suffix in candidates: url = f"{API_BASE}{suffix}" try: r = await client.get(url, timeout=60, follow_redirects=True) if r.status_code == 200: # 下载 zip 文件 zip_path = os.path.join(output_dir, f"{task_id}.zip") with open(zip_path, "wb") as f: f.write(r.content) return zip_path except Exception: continue raise RuntimeError(f"Failed to download result for task {task_id}") def extract_readme(zip_path: str, output_dir: str) -> str: """解压结果并返回 markdown 内容""" extract_dir = os.path.join(output_dir, "extracted") os.makedirs(extract_dir, exist_ok=True) with zipfile.ZipFile(zip_path, "r") as zf: zf.extractall(extract_dir) # 查找 markdown 文件 md_files = list(Path(extract_dir).rglob("*.md")) if md_files: return md_files[0].read_text(encoding="utf-8") return "No markdown output found in result." # --------------------------------------------------------------------------- # Gradio 处理函数 # --------------------------------------------------------------------------- async def process_pdf(file_obj, backend, lang, progress=gr.Progress()): """Gradio 事件处理:上传 → 提交 → 轮询 → 返回结果""" if file_obj is None: return "Please upload a PDF file.", "", "" tmp_dir = tempfile.mkdtemp(prefix="mineru-client-") file_path = file_obj.name if hasattr(file_obj, "name") else os.path.join(tmp_dir, "input.pdf") file_name = os.path.basename(file_path) async with httpx.AsyncClient(timeout=httpx.Timeout(120)) as client: # 1. 探测 API progress(0.05, desc="Connecting to API...") api_prefix = await discover_api(client) progress(0.1, desc=f"API prefix: {api_prefix}") # 2. 提交任务 progress(0.15, desc="Submitting task...") try: submit_resp = await submit_task(client, api_prefix, file_path, file_name, backend, lang) except httpx.HTTPStatusError as e: # 尝试旧版 endpoint(不带 /submit 后缀) alt_data = {"backend": backend, "lang": lang} try: alt_url = f"{API_BASE}{api_prefix}" with open(file_path, "rb") as f: files = {"file": (file_name, f, "application/pdf")} r = await client.post(alt_url, files=files, data=alt_data, timeout=30) r.raise_for_status() submit_resp = r.json() except Exception as e2: return f"Task submission failed: {e}\nAlt attempt: {e2}", "", "" task_id = submit_resp.get("task_id") or submit_resp.get("data", {}).get("task_id") if not task_id: return f"Unexpected submit response:\n{json.dumps(submit_resp, indent=2)}", "", "" progress(0.2, desc=f"Task ID: {task_id}") # 3. 轮询状态 status_text = "" start_time = time.time() last_idx = -1 tickers = ["⣾", "⣽", "⣻", "⢿", "⡿", "⣟", "⣯", "⣷"] while True: elapsed = time.time() - start_time if elapsed > MAX_WAIT: return f"Task timed out after {MAX_WAIT}s\nLast status:\n{status_text}", "", "" t = (int(elapsed * 2)) % len(tickers) progress(min(0.2 + 0.6 * (elapsed / 60), 0.8), desc=f"{tickers[t]} Processing... ({int(elapsed)}s)") try: status_resp = await get_task_status(client, api_prefix, task_id) except Exception as e: status_text = f"Polling error: {e}" await asyncio.sleep(POLL_INTERVAL) continue status = status_resp.get("status", "unknown") progress_pct = status_resp.get("progress", 0) err_msg = status_resp.get("error", "") status_text = json.dumps(status_resp, indent=2, ensure_ascii=False) if status == "completed" or status == "success": progress(0.85, desc="Downloading result...") break elif status == "failed": return f"## Task Failed\n\nError: {err_msg}\n\n```json\n{status_text}\n```", "", "" elif status == "processing": pass else: pass await asyncio.sleep(POLL_INTERVAL) # 4. 下载结果 progress(0.9, desc="Downloading result...") try: zip_path = await download_result(client, api_prefix, task_id, tmp_dir) except Exception as e: return f"Result download failed: {e}\n\nStatus:\n{status_text}", "", "" # 5. 提取 markdown progress(0.95, desc="Extracting results...") try: md_content = extract_readme(zip_path, tmp_dir) except Exception as e: md_content = f"Extraction error: {e}" progress(1.0, desc="Done!") return md_content, status_text, zip_path # --------------------------------------------------------------------------- # Gradio UI # --------------------------------------------------------------------------- def create_ui(): custom_css = """ .result-panel { min-height: 400px; } .status-panel { min-height: 200px; font-size: 13px; } """ with gr.Blocks(css=custom_css, title="MinerU ROCm — Document Parser") as demo: gr.Markdown(""" # MinerU ROCm — Document Parser Upload a PDF file. Processing is handled by the **mineru-router → mineru-worker** backend. """) with gr.Row(): with gr.Column(scale=1): file_input = gr.File( label="Upload PDF", file_types=[".pdf"], type="filepath", ) backend_dd = gr.Dropdown( label="Backend", choices=["hybrid-auto-engine", "hybrid-http-client", "vlm-http-client", "pipeline", "auto"], value=DEFAULT_BACKEND, ) lang_dd = gr.Dropdown( label="Language", choices=["ch", "en", "japan", "korean"], value=DEFAULT_LANG, ) submit_btn = gr.Button("Parse Document", variant="primary", size="lg") status_display = gr.Textbox( label="Status", lines=6, max_lines=12, elem_classes=["status-panel"], ) zip_display = gr.File( label="Download Result ZIP", type="filepath", visible=True, ) with gr.Column(scale=2): result_display = gr.Markdown( value="*Upload a PDF to start parsing...*", elem_classes=["result-panel"], ) submit_btn.click( fn=process_pdf, inputs=[file_input, backend_dd, lang_dd], outputs=[result_display, status_display, zip_display], ) return demo if __name__ == "__main__": server_name = os.environ.get("GRADIO_SERVER_NAME", "0.0.0.0") server_port = int(os.environ.get("GRADIO_SERVER_PORT", "7860")) demo = create_ui() demo.queue(default_concurrency_limit=3, max_size=10).launch( server_name=server_name, server_port=server_port, share=False, )