262 lines
9.1 KiB
Python
262 lines
9.1 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 版本的 file 返回值
|
|
if isinstance(file_obj, str):
|
|
file_path = file_obj
|
|
file_name = os.path.basename(file_path)
|
|
elif hasattr(file_obj, 'name'):
|
|
file_path = file_obj.name
|
|
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.File(
|
|
label="Upload PDF",
|
|
file_types=[".pdf"],
|
|
type="filepath",
|
|
)
|
|
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"],
|
|
)
|
|
|
|
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,
|
|
css=CUSTOM_CSS,
|
|
)
|