devappsmi/document_parse
0
1"""2 3PaddleOCR-VL-1.5 Bridge Server (HF Spaces Edition)4====================================================5 6With per-token and per-word confidence scores via vLLM logprobs.7 8Architecture:9 Gradio App → This Bridge (port 7860) → vLLM Docker (117.54.141.62:8000)10"""11 12import base6413import json14import math15import os16import shutil17import tempfile18import traceback19import uuid20from typing import Any, Dict, List, Optional, Tuple21 22import uvicorn23from fastapi import FastAPI, File, Header, HTTPException, Request, UploadFile24from fastapi.middleware.cors import CORSMiddleware25from fastapi.staticfiles import StaticFiles26from openai import OpenAI27from PIL import Image28 29# =============================================================================30# Configuration31# =============================================================================32VLLM_SERVER_URL = os.environ.get("VLLM_SERVER_URL", "http://117.54.141.62:8000/v1")33VLLM_MODEL_NAME = os.environ.get("VLLM_MODEL_NAME", "PaddleOCR-VL-1.5-0.9B")34BRIDGE_PORT = int(os.environ.get("PORT", "7860"))35API_KEY = os.environ.get("API_KEY", "")36 37SPACE_HOST = os.environ.get("SPACE_HOST", "")38if SPACE_HOST:39 PUBLIC_BASE_URL = f"https://{SPACE_HOST}"40else:41 PUBLIC_BASE_URL = os.environ.get("PUBLIC_BASE_URL", f"http://localhost:{BRIDGE_PORT}")42 43STATIC_DIR = "/tmp/ocr_outputs"44os.makedirs(STATIC_DIR, exist_ok=True)45 46# =============================================================================47# Initialize clients48# =============================================================================49openai_client = OpenAI(50 api_key="EMPTY",51 base_url=VLLM_SERVER_URL,52 timeout=60053)54 55pipeline = None56 57 58def get_pipeline():59 global pipeline60 if pipeline is None:61 from paddleocr import PaddleOCRVL62 pipeline = PaddleOCRVL(63 vl_rec_backend="vllm-server",64 vl_rec_server_url=VLLM_SERVER_URL65 )66 return pipeline67 68 69# =============================================================================70# FastAPI App71# =============================================================================72app = FastAPI(73 title="PaddleOCR-VL-1.5 Bridge API",74 description="Full document parsing API with per-token/word confidence scores",75 version="1.1.0"76)77 78app.add_middleware(79 CORSMiddleware,80 allow_origins=["*"],81 allow_credentials=True,82 allow_methods=["*"],83 allow_headers=["*"],84)85 86app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")87 88 89# =============================================================================90# Auth91# =============================================================================92def verify_auth(authorization: Optional[str] = None):93 if API_KEY and API_KEY.strip():94 if not authorization or authorization != f"Bearer {API_KEY}":95 raise HTTPException(status_code=401, detail="Unauthorized")96 97 98# =============================================================================99# Confidence Score Helpers100# =============================================================================101 102def parse_logprobs(response) -> List[Dict[str, Any]]:103 """104 Extract per-token confidence from the OpenAI response logprobs.105 Returns list of {token, logprob, confidence} dicts.106 """107 token_details = []108 109 try:110 choice = response.choices[0]111 logprobs_data = choice.logprobs112 113 if logprobs_data is None:114 return token_details115 116 # OpenAI format: logprobs.content is a list of token info117 content_logprobs = getattr(logprobs_data, 'content', None)118 119 if content_logprobs:120 # OpenAI-compatible format (newer vLLM)121 for token_info in content_logprobs:122 token_str = getattr(token_info, 'token', '')123 logprob_val = getattr(token_info, 'logprob', None)124 125 if logprob_val is not None:126 confidence = math.exp(logprob_val) # convert log prob to probability127 else:128 confidence = 0.0129 logprob_val = float('-inf')130 131 token_details.append({132 "token": token_str,133 "logprob": round(logprob_val, 6),134 "confidence": round(confidence, 6)135 })136 else:137 # Legacy vLLM format: logprobs has tokens, token_logprobs138 tokens = getattr(logprobs_data, 'tokens', None)139 token_logprobs = getattr(logprobs_data, 'token_logprobs', None)140 141 if tokens and token_logprobs:142 for token_str, logprob_val in zip(tokens, token_logprobs):143 if logprob_val is not None:144 confidence = math.exp(logprob_val)145 else:146 confidence = 0.0147 logprob_val = float('-inf')148 149 token_details.append({150 "token": token_str,151 "logprob": round(logprob_val, 6),152 "confidence": round(confidence, 6)153 })154 155 except Exception as e:156 print(f"Warning: Could not parse logprobs: {e}")157 traceback.print_exc()158 159 return token_details160 161 162def tokens_to_words(token_details: List[Dict[str, Any]]) -> List[Dict[str, Any]]:163 """164 Group tokens into words. A new word starts when a token begins with a space165 or is a newline. Returns list of {word, tokens, confidence, avg_logprob}.166 167 Word confidence = geometric mean of token probabilities168 = exp(mean of logprobs)169 """170 if not token_details:171 return []172 173 words = []174 current_word_tokens = []175 176 for td in token_details:177 token = td["token"]178 179 # Detect word boundary: starts with space, is newline, or is punctuation-only after text180 is_boundary = (181 token.startswith(" ") or182 token.startswith("▁") or # sentencepiece space marker183 token.startswith("Ġ") or # GPT-2 style space marker184 token in ("\n", "\r", "\t", "\r\n") or185 (len(current_word_tokens) > 0 and token.strip() == "")186 )187 188 if is_boundary and current_word_tokens:189 # Finalize previous word190 words.append(_finalize_word(current_word_tokens))191 current_word_tokens = []192 193 current_word_tokens.append(td)194 195 # Don't forget the last word196 if current_word_tokens:197 words.append(_finalize_word(current_word_tokens))198 199 return words200 201 202def _finalize_word(tokens: List[Dict[str, Any]]) -> Dict[str, Any]:203 """Compute word-level confidence from its constituent tokens."""204 # Reconstruct word text205 word_text = "".join(t["token"] for t in tokens).strip()206 # Remove sentencepiece/GPT markers207 word_text = word_text.lstrip("▁Ġ ")208 209 # Geometric mean of probabilities = exp(mean of logprobs)210 valid_logprobs = [t["logprob"] for t in tokens if t["logprob"] != float('-inf')]211 if valid_logprobs:212 avg_logprob = sum(valid_logprobs) / len(valid_logprobs)213 word_confidence = math.exp(avg_logprob)214 else:215 avg_logprob = float('-inf')216 word_confidence = 0.0217 218 return {219 "word": word_text,220 "confidence": round(word_confidence, 6),221 "avg_logprob": round(avg_logprob, 6) if avg_logprob != float('-inf') else None,222 "token_count": len(tokens),223 "tokens": [224 {"token": t["token"], "confidence": t["confidence"]}225 for t in tokens226 ]227 }228 229 230def compute_overall_confidence(token_details: List[Dict[str, Any]]) -> Dict[str, Any]:231 """Compute overall text confidence statistics."""232 if not token_details:233 return {"mean_confidence": 0.0, "min_confidence": 0.0, "total_tokens": 0}234 235 confidences = [t["confidence"] for t in token_details]236 logprobs = [t["logprob"] for t in token_details if t["logprob"] != float('-inf')]237 238 mean_conf = sum(confidences) / len(confidences) if confidences else 0.0239 min_conf = min(confidences) if confidences else 0.0240 max_conf = max(confidences) if confidences else 0.0241 242 # Perplexity = exp(-mean(logprobs)) — lower is more confident243 if logprobs:244 avg_logprob = sum(logprobs) / len(logprobs)245 perplexity = math.exp(-avg_logprob)246 else:247 perplexity = float('inf')248 249 return {250 "mean_confidence": round(mean_conf, 6),251 "min_confidence": round(min_conf, 6),252 "max_confidence": round(max_conf, 6),253 "perplexity": round(perplexity, 4) if perplexity != float('inf') else None,254 "total_tokens": len(token_details)255 }256 257 258# =============================================================================259# Image / File Helpers260# =============================================================================261TASK_PROMPTS = {262 "ocr": "OCR:",263 "formula": "Formula Recognition:",264 "table": "Table Recognition:",265 "chart": "Chart Recognition:",266 "spotting": "Spotting:",267 "seal": "Seal Recognition:",268}269 270IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp", ".bmp", ".gif"}271 272 273def save_temp_image(file_data: str) -> str:274 if file_data.startswith(("http://", "https://")):275 import requests as req276 resp = req.get(file_data, timeout=120)277 resp.raise_for_status()278 content = resp.content279 ct = resp.headers.get("content-type", "image/png")280 ext = ".png"281 if "jpeg" in ct or "jpg" in ct:282 ext = ".jpg"283 elif "webp" in ct:284 ext = ".webp"285 elif "bmp" in ct:286 ext = ".bmp"287 else:288 content = base64.b64decode(file_data)289 ext = ".png"290 291 tmp = tempfile.NamedTemporaryFile(delete=False, suffix=ext)292 tmp.write(content)293 tmp.close()294 return tmp.name295 296 297def serve_file(src_path: str, request_id: str, filename: str) -> str:298 static_subdir = os.path.join(STATIC_DIR, request_id)299 os.makedirs(static_subdir, exist_ok=True)300 dst_path = os.path.join(static_subdir, filename)301 shutil.copy2(src_path, dst_path)302 return f"{PUBLIC_BASE_URL}/static/{request_id}/{filename}"303 304 305def collect_images_from_dir(directory: str, request_id: str) -> Dict[str, str]:306 result = {}307 if not os.path.exists(directory):308 return result309 for root, dirs, files in os.walk(directory):310 for fname in files:311 ext = os.path.splitext(fname)[1].lower()312 if ext in IMAGE_EXTENSIONS:313 src = os.path.join(root, fname)314 rel_path = os.path.relpath(src, directory)315 safe_name = rel_path.replace(os.sep, "_")316 url = serve_file(src, request_id, safe_name)317 result[rel_path] = url318 return result319 320 321# =============================================================================322# VLM call with confidence323# =============================================================================324 325def call_vllm_with_confidence(image_url: str, task_prompt: str) -> Tuple[str, List[Dict], List[Dict], Dict]:326 """327 Call vLLM with logprobs enabled.328 Returns: (result_text, token_confidences, word_confidences, overall_stats)329 """330 response = openai_client.chat.completions.create(331 model=VLLM_MODEL_NAME,332 messages=[{333 "role": "user",334 "content": [335 {"type": "image_url", "image_url": {"url": image_url}},336 {"type": "text", "text": task_prompt}337 ]338 }],339 temperature=0.0,340 logprobs=True,341 top_logprobs=5342 )343 344 result_text = response.choices[0].message.content345 346 # Extract per-token confidence347 token_details = parse_logprobs(response)348 349 # Group into words350 word_details = tokens_to_words(token_details)351 352 # Overall stats353 overall_stats = compute_overall_confidence(token_details)354 355 return result_text, token_details, word_details, overall_stats356 357 358# =============================================================================359# Element-level Recognition360# =============================================================================361 362def element_level_recognition(file_data: str, prompt_label: str) -> Dict[str, Any]:363 """Element-level recognition with confidence scores."""364 if file_data.startswith(("http://", "https://")):365 image_url = file_data366 else:367 image_url = f"data:image/png;base64,{file_data}"368 369 task_prompt = TASK_PROMPTS.get(prompt_label, "OCR:")370 371 result_text, token_details, word_details, overall_stats = call_vllm_with_confidence(372 image_url, task_prompt373 )374 375 return {376 "errorCode": 0,377 "result": {378 "layoutParsingResults": [{379 "prunedResult": {380 "page_count": 1,381 "width": 0,382 "height": 0,383 "parsing_res_list": [{384 "block_label": prompt_label,385 "block_content": result_text,386 "block_bbox": [],387 "block_id": 0,388 "block_order": 0,389 "group_id": 0,390 "global_block_id": 0,391 "global_group_id": 0,392 "block_polygon_points": []393 }],394 "layout_det_res": {"boxes": []},395 "spotting_res": _parse_spotting(result_text) if prompt_label == "spotting" else {}396 },397 "markdown": {"text": result_text, "images": {}},398 "outputImages": {},399 "confidence": {400 "overall": overall_stats,401 "tokens": token_details,402 "words": word_details403 }404 }]405 }406 }407 408 409# =============================================================================410# Full Document Parsing411# =============================================================================412 413def full_document_parsing(file_data: str, use_chart_recognition: bool = False,414 use_doc_unwarping: bool = True,415 use_doc_orientation_classify: bool = True,416 include_confidence: bool = True) -> Dict[str, Any]:417 """418 Full document parsing with layout detection + VLM recognition.419 When include_confidence=True, re-runs each block through vLLM with logprobs420 to get per-token/word confidence scores.421 """422 tmp_path = save_temp_image(file_data)423 request_id = str(uuid.uuid4())[:12]424 425 try:426 # Get image dimensions427 try:428 img = Image.open(tmp_path)429 img_width, img_height = img.size430 img.close()431 except Exception:432 img_width, img_height = 0, 0433 434 pipe = get_pipeline()435 output = pipe.predict(tmp_path)436 437 layout_parsing_results = []438 preprocessed_images = []439 data_info_pages = []440 441 for i, res in enumerate(output):442 page_id = f"{request_id}_p{i}"443 output_dir = tempfile.mkdtemp()444 445 # Save all outputs446 res.save_to_json(save_path=output_dir)447 res.save_to_markdown(save_path=output_dir)448 try:449 res.save_to_img(save_path=output_dir)450 except Exception:451 pass452 453 # --- Read markdown ---454 md_text = ""455 md_files = [f for f in os.listdir(output_dir) if f.endswith(".md")]456 if md_files:457 with open(os.path.join(output_dir, md_files[0]), "r", encoding="utf-8") as f:458 md_text = f.read()459 460 # --- Read JSON ---461 json_data = {}462 json_files = [f for f in os.listdir(output_dir) if f.endswith(".json")]463 if json_files:464 with open(os.path.join(output_dir, json_files[0]), "r", encoding="utf-8") as f:465 json_data = json.load(f)466 467 # --- Collect and serve images ---468 all_images = collect_images_from_dir(output_dir, page_id)469 470 output_images = {}471 for rel_path, url in all_images.items():472 name = os.path.splitext(os.path.basename(rel_path))[0]473 if "layout" in name.lower() or "det" in name.lower() or "vis" in name.lower():474 output_images["layout_det_res"] = url475 else:476 output_images[name] = url477 478 md_images = {}479 imgs_dir = os.path.join(output_dir, "imgs")480 if os.path.exists(imgs_dir):481 for fname in os.listdir(imgs_dir):482 ext = os.path.splitext(fname)[1].lower()483 if ext in IMAGE_EXTENSIONS:484 src = os.path.join(imgs_dir, fname)485 url = serve_file(src, page_id, fname)486 local_ref = f"imgs/{fname}"487 md_images[local_ref] = url488 md_text = md_text.replace(f'src="{local_ref}"', f'src="{url}"')489 md_text = md_text.replace(f']({local_ref})', f']({url})')490 491 input_image_url = serve_file(tmp_path, page_id, f"input_img_{i}.jpg")492 493 # --- Build prunedResult ---494 pruned_result = {}495 if json_data:496 pruned_result = {497 "page_count": json_data.get("page_count", 1),498 "width": json_data.get("width", img_width),499 "height": json_data.get("height", img_height),500 "model_settings": json_data.get("model_settings", {501 "use_doc_preprocessor": False,502 "use_layout_detection": True,503 "use_chart_recognition": use_chart_recognition,504 "use_seal_recognition": True,505 "use_ocr_for_image_block": False,506 "format_block_content": True,507 "merge_layout_blocks": True,508 "markdown_ignore_labels": [509 "number", "footnote", "header",510 "header_image", "footer", "footer_image", "aside_text"511 ],512 "return_layout_polygon_points": True513 }),514 "parsing_res_list": json_data.get("parsing_res_list",515 json_data.get("blocks", [])),516 "layout_det_res": json_data.get("layout_det_res",517 json_data.get("det_res", {"boxes": []}))518 }519 else:520 pruned_result = {521 "page_count": 1,522 "width": img_width,523 "height": img_height,524 "model_settings": {},525 "parsing_res_list": [],526 "layout_det_res": {"boxes": []}527 }528 529 if not pruned_result.get("width"):530 pruned_result["width"] = img_width531 if not pruned_result.get("height"):532 pruned_result["height"] = img_height533 534 # --- Confidence scores for each block ---535 block_confidences = []536 if include_confidence and pruned_result.get("parsing_res_list"):537 # Use the full-page image for confidence scoring538 if file_data.startswith(("http://", "https://")):539 conf_image_url = file_data540 else:541 conf_image_url = f"data:image/png;base64,{file_data}"542 543 # Get confidence for the entire page text544 try:545 _, page_tokens, page_words, page_overall = call_vllm_with_confidence(546 conf_image_url, "OCR:"547 )548 block_confidences = {549 "overall": page_overall,550 "tokens": page_tokens,551 "words": page_words552 }553 except Exception as e:554 print(f"Warning: Could not get confidence scores: {e}")555 block_confidences = {556 "overall": {"mean_confidence": 0, "total_tokens": 0},557 "tokens": [],558 "words": []559 }560 561 # --- Build page result ---562 page_result = {563 "prunedResult": pruned_result,564 "markdown": {565 "text": md_text,566 "images": md_images567 },568 "outputImages": output_images,569 "inputImage": input_image_url,570 }571 572 if block_confidences:573 page_result["confidence"] = block_confidences574 575 layout_parsing_results.append(page_result)576 preprocessed_images.append(input_image_url)577 data_info_pages.append({578 "width": img_width,579 "height": img_height580 })581 582 return {583 "errorCode": 0,584 "result": {585 "layoutParsingResults": layout_parsing_results if layout_parsing_results else [{586 "prunedResult": {587 "page_count": 0, "width": 0, "height": 0,588 "parsing_res_list": [], "layout_det_res": {"boxes": []}589 },590 "markdown": {"text": "", "images": {}},591 "outputImages": {},592 "inputImage": ""593 }],594 "preprocessedImages": preprocessed_images,595 "dataInfo": {596 "type": "image",597 "numPages": len(layout_parsing_results),598 "pages": data_info_pages599 }600 }601 }602 603 finally:604 if os.path.exists(tmp_path):605 os.unlink(tmp_path)606 607 608def _parse_spotting(text: str) -> dict:609 try:610 return json.loads(text)611 except (json.JSONDecodeError, TypeError):612 return {"raw_text": text}613 614 615# =============================================================================616# Endpoints617# =============================================================================618 619@app.get("/")620async def root():621 return {622 "service": "PaddleOCR-VL-1.5 Bridge API",623 "status": "running",624 "version": "1.1.0 (with confidence scores)",625 "endpoints": ["/health", "/api/ocr", "/api/parse", "/api/parse/markdown", "/v1/chat/completions", "/docs"]626 }627 628 629@app.get("/health")630async def health():631 return {"status": "ok", "model": VLLM_MODEL_NAME, "vllm_url": VLLM_SERVER_URL}632 633 634@app.post("/api/ocr")635async def ocr_endpoint(request: Request, authorization: Optional[str] = Header(None)):636 """637 Main OCR endpoint — compatible with the Gradio app.638 Now includes per-token and per-word confidence scores.639 640 Body:641 {642 "file": "base64_or_url",643 "useLayoutDetection": true/false,644 "promptLabel": "ocr|formula|table|chart|spotting|seal",645 "useChartRecognition": false,646 "useDocUnwarping": true,647 "useDocOrientationClassify": true,648 "includeConfidence": true (default: true)649 }650 651 Response includes:652 {653 "result": {654 "layoutParsingResults": [{655 ...656 "confidence": {657 "overall": {658 "mean_confidence": 0.95,659 "min_confidence": 0.42,660 "max_confidence": 1.0,661 "perplexity": 1.12,662 "total_tokens": 85663 },664 "tokens": [665 {"token": "Hello", "logprob": -0.02, "confidence": 0.98},666 ...667 ],668 "words": [669 {"word": "Hello", "confidence": 0.98, "avg_logprob": -0.02, "token_count": 1, "tokens": [...]},670 ...671 ]672 }673 }]674 }675 }676 """677 verify_auth(authorization)678 679 try:680 body = await request.json()681 except Exception:682 raise HTTPException(status_code=400, detail="Invalid JSON body")683 684 file_data = body.get("file", "")685 if not file_data:686 raise HTTPException(status_code=400, detail="Missing 'file' field")687 688 use_layout = body.get("useLayoutDetection", False)689 prompt_label = body.get("promptLabel", "ocr")690 use_chart = body.get("useChartRecognition", False)691 use_unwarp = body.get("useDocUnwarping", True)692 use_orient = body.get("useDocOrientationClassify", True)693 include_confidence = body.get("includeConfidence", True)694 695 try:696 if use_layout:697 return full_document_parsing(698 file_data, use_chart, use_unwarp, use_orient,699 include_confidence=include_confidence700 )701 else:702 return element_level_recognition(file_data, prompt_label)703 except Exception as e:704 traceback.print_exc()705 return {"errorCode": -1, "errorMsg": str(e)}706 707 708@app.post("/api/parse")709async def parse_file(710 file: UploadFile = File(...),711 use_layout_detection: bool = True,712 prompt_label: str = "ocr",713 include_confidence: bool = True,714 authorization: Optional[str] = Header(None)715):716 """File upload endpoint with confidence scores."""717 verify_auth(authorization)718 content = await file.read()719 b64 = base64.b64encode(content).decode("utf-8")720 721 try:722 if use_layout_detection:723 return full_document_parsing(b64, include_confidence=include_confidence)724 else:725 return element_level_recognition(b64, prompt_label)726 except Exception as e:727 traceback.print_exc()728 return {"errorCode": -1, "errorMsg": str(e)}729 730 731@app.post("/api/parse/markdown")732async def parse_to_markdown(733 file: UploadFile = File(...),734 authorization: Optional[str] = Header(None)735):736 """Returns just markdown text."""737 verify_auth(authorization)738 content = await file.read()739 b64 = base64.b64encode(content).decode("utf-8")740 741 try:742 result = full_document_parsing(b64, include_confidence=False)743 pages = result.get("result", {}).get("layoutParsingResults", [])744 markdown_parts = [p.get("markdown", {}).get("text", "") for p in pages if p.get("markdown", {}).get("text")]745 return {746 "status": "ok",747 "markdown": "\n\n---\n\n".join(markdown_parts),748 "page_count": len(pages)749 }750 except Exception as e:751 traceback.print_exc()752 raise HTTPException(status_code=500, detail=str(e))753 754 755@app.post("/v1/chat/completions")756async def proxy_chat_completions(request: Request, authorization: Optional[str] = Header(None)):757 """Proxy to vLLM for direct OpenAI-compatible calls (logprobs supported)."""758 verify_auth(authorization)759 760 import httpx761 body = await request.json()762 763 async with httpx.AsyncClient(timeout=600) as client:764 resp = await client.post(765 f"{VLLM_SERVER_URL}/chat/completions",766 json=body,767 headers={"Content-Type": "application/json"}768 )769 return resp.json()770 771 772# =============================================================================773# Entry point774# =============================================================================775if __name__ == "__main__":776 print(f"""777╔══════════════════════════════════════════════════════════════╗778║ PaddleOCR-VL-1.5 Bridge Server (HF Spaces) ║779║ v1.1.0 — with per-token/word confidence scores ║780╠══════════════════════════════════════════════════════════════╣781║ Bridge API: http://0.0.0.0:{BRIDGE_PORT} ║782║ vLLM backend: {VLLM_SERVER_URL:<44s}║783║ Model: {VLLM_MODEL_NAME:<44s}║784║ Auth: {"ENABLED" if API_KEY else "DISABLED":<44s}║785╠══════════════════════════════════════════════════════════════╣786║ Endpoints: ║787║ GET /health - Health check ║788║ GET /docs - Swagger UI ║789║ POST /api/ocr - Gradio-compatible + confidence║790║ POST /api/parse - File upload + confidence ║791║ POST /api/parse/markdown - Simple markdown output ║792║ POST /v1/chat/completions - vLLM proxy (OpenAI format) ║793║ GET /static/... - Output images ║794╚══════════════════════════════════════════════════════════════╝795 """)796 uvicorn.run(app, host="0.0.0.0", port=BRIDGE_PORT)