cjw/mineru-rocm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

x

Browse Source
This commit is contained in:
chenjw28
2026-06-15 17:46:45 +08:00
parent 2b779d0423
commit c33488d95c
3 changed files with 56 additions and 74 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docker/scripts/gradio_client.py
+23 -47
View File
@@ -17,44 +17,40 @@ 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")
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"))
async def discover_api(client: httpx.AsyncClient) -> str:
"""探测实际可用的 API 前缀: /api/v1/tasks 或 /tasks"""
"""探测实际可用的 API 前缀: /file_parse 端点"""
candidates = [
f"{API_BASE}/api/v1/health",
f"{API_BASE}/health",
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 < 500:
# FastAPI docs → /api/v1 ; plain health → /
if "/api/v1" in str(r.url):
return "/api/v1/tasks"
return "/tasks"
if r.status_code == 200:
return "" # worker uses root-level endpoints like /file_parse
except Exception:
continue
# 默认新版 API
return "/api/v1/tasks"
return ""
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"
"""提交解析任务到 worker 的 /file_parse 端点"""
submit_url = f"{API_BASE}/file_parse"
with open(file_path, "rb") as f:
files = {"file": (file_name, f, "application/pdf")}
files = {"files": (file_name, f, "application/pdf")}
data = {
"backend": backend,
"lang": lang,
"lang_list": lang,
"parse_method": "auto",
}
r = await client.post(submit_url, files=files, data=data, timeout=30)
r.raise_for_status()
@@ -64,7 +60,7 @@ async def submit_task(client: httpx.AsyncClient, api_prefix: str,
async def get_task_status(client: httpx.AsyncClient, api_prefix: str,
task_id: str) -> dict:
"""查询任务状态"""
url = f"{API_BASE}{api_prefix}/{task_id}"
url = f"{API_BASE}/tasks/{task_id}"
r = await client.get(url, timeout=10)
r.raise_for_status()
return r.json()
@@ -72,24 +68,14 @@ async def get_task_status(client: httpx.AsyncClient, api_prefix: str,
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}")
"""下载任务结果"""
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:
@@ -127,18 +113,8 @@ async def process_pdf(file_obj, backend, lang, progress=gr.Progress()):
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}", "", ""
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: