269 lines
9.4 KiB
Python
269 lines
9.4 KiB
Python
#!/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-nginx:8000")
|
||
DEFAULT_BACKEND = os.environ.get("MINERU_BACKEND", "hybrid-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"))
|
||
|
||
CUSTOM_CSS = """
|
||
.result-panel { min-height: 400px; }
|
||
.status-panel { min-height: 200px; font-size: 13px; }
|
||
"""
|
||
|
||
|
||
async def discover_api(client: httpx.AsyncClient) -> str:
|
||
"""探测实际可用的 API 前缀: /file_parse 端点"""
|
||
candidates = [
|
||
f"{API_BASE}/openapi.json",
|
||
f"{API_BASE}/docs",
|
||
]
|
||
for url in candidates:
|
||
try:
|
||
r = await client.get(url, timeout=5)
|
||
if r.status_code == 200:
|
||
return "" # worker uses root-level endpoints like /file_parse
|
||
except Exception:
|
||
continue
|
||
return ""
|
||
|
||
|
||
async def submit_task(client: httpx.AsyncClient, api_prefix: str,
|
||
file_path: str, file_name: str,
|
||
backend: str, lang: str) -> dict:
|
||
"""提交解析任务到 worker 的 /file_parse 端点"""
|
||
submit_url = f"{API_BASE}/file_parse"
|
||
with open(file_path, "rb") as f:
|
||
files = {"files": (file_name, f, "application/pdf")}
|
||
data = {
|
||
"backend": backend,
|
||
"lang_list": lang,
|
||
"parse_method": "auto",
|
||
}
|
||
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}/tasks/{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:
|
||
"""下载任务结果"""
|
||
url = f"{API_BASE}/tasks/{task_id}/result"
|
||
r = await client.get(url, timeout=60, follow_redirects=True)
|
||
r.raise_for_status()
|
||
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
|
||
|
||
|
||
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-")
|
||
|
||
# Gradio 6.0 binary mode：file_obj 是 bytes
|
||
if isinstance(file_obj, bytes):
|
||
file_path = os.path.join(tmp_dir, "input.pdf")
|
||
file_name = "input.pdf"
|
||
with open(file_path, "wb") as f:
|
||
f.write(file_obj)
|
||
elif isinstance(file_obj, str):
|
||
file_path = file_obj
|
||
file_name = os.path.basename(file_path)
|
||
else:
|
||
return "Unsupported file object type", "", ""
|
||
|
||
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 Exception as e:
|
||
return f"Task submission failed: {e}", "", ""
|
||
|
||
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():
|
||
with gr.Blocks(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.UploadButton(
|
||
label="Upload PDF",
|
||
file_types=[".pdf"],
|
||
type="binary",
|
||
)
|
||
backend_dd = gr.Dropdown(
|
||
label="Backend",
|
||
choices=["hybrid-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"],
|
||
)
|
||
|
||
file_input.upload(
|
||
fn=process_pdf,
|
||
inputs=[file_input, backend_dd, lang_dd],
|
||
outputs=[result_display, status_display, zip_display],
|
||
)
|
||
submit_btn.click(
|
||
fn=lambda f,b,l: "Please use the Upload button above.",
|
||
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,
|
||
css=CUSTOM_CSS,
|
||
)
|