Spaces:
Running on Zero
Running on Zero
| """Gradio demo for ATH-MaaS/OvisOCR3 — a unified 0.8B OCR / document-understanding | |
| vision-language model covering 5 capability domains selected via task prompts: | |
| document parsing, general OCR, key information extraction, figure captioning, | |
| and cross-page table reconstruction. | |
| The inference path follows the model card exactly (task prompts from the repo's | |
| `prompts/` directory, literal `.replace()` substitution, greedy decoding, | |
| `enable_thinking=False`, 448²–2880² pixel budget) and is served through | |
| transformers' `Qwen3_5ForConditionalGeneration` on ZeroGPU. | |
| """ | |
| import json | |
| import os | |
| import re | |
| from pathlib import Path | |
| os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") | |
| import spaces # noqa: E402 (must precede torch / CUDA-touching imports) | |
| import gradio as gr # noqa: E402 | |
| import torch # noqa: E402 | |
| from transformers import AutoProcessor, Qwen3_5ForConditionalGeneration # noqa: E402 | |
| MODEL_ID = "ATH-MaaS/OvisOCR3" | |
| PROMPTS_DIR = Path(__file__).parent / "prompts" | |
| # Pixel budget from the model card's `images_kwargs`. | |
| MIN_PIXELS = 448 * 448 | |
| MAX_PIXELS = 2880 * 2880 | |
| MAX_NEW_TOKENS_CEILING = 16384 | |
| # --------------------------------------------------------------------------- | |
| # Model + processor, loaded eagerly at module scope on CUDA (ZeroGPU pattern). | |
| # --------------------------------------------------------------------------- | |
| processor = AutoProcessor.from_pretrained( | |
| MODEL_ID, | |
| min_pixels=MIN_PIXELS, | |
| max_pixels=MAX_PIXELS, | |
| ) | |
| model = Qwen3_5ForConditionalGeneration.from_pretrained( | |
| MODEL_ID, | |
| dtype=torch.bfloat16, | |
| attn_implementation="sdpa", | |
| ).to("cuda") | |
| model.eval() | |
| # --------------------------------------------------------------------------- | |
| # Prompt loading — mirrors the model card's `load_prompt`. | |
| # --------------------------------------------------------------------------- | |
| def load_prompt(name: str) -> str: | |
| """Read a task prompt shipped with the Space from `prompts/<name>.txt`. | |
| Args: | |
| name: Prompt file stem, e.g. "document_markdown". | |
| Returns: | |
| The raw prompt text exactly as shipped in the model repo. | |
| """ | |
| return (PROMPTS_DIR / f"{name}.txt").read_text(encoding="utf-8") | |
| def _general_ocr_prompt(name: str) -> str: | |
| """Load a General OCR task prompt, stripped except for position-aware VQA. | |
| Args: | |
| name: Task prompt stem (see model card's `general_ocr_prompt`). | |
| Returns: | |
| The task prompt, stripped unless it is `position_aware_vqa`. | |
| """ | |
| prompt = load_prompt(name) | |
| return prompt if name == "position_aware_vqa" else prompt.strip() | |
| # --------------------------------------------------------------------------- | |
| # Post-processing helpers — copied from the model card. | |
| # --------------------------------------------------------------------------- | |
| def _clean_truncated_repeats( | |
| text: str, | |
| min_text_len: int = 8000, | |
| max_period: int = 200, | |
| min_period: int = 1, | |
| min_repeat_chars: int = 100, | |
| min_repeat_times: int = 5, | |
| ) -> str: | |
| """Remove truncated repetition artifacts from generated text (model card).""" | |
| n = len(text) | |
| if n < min_text_len: | |
| return text | |
| max_period = min(max_period, n - 1) | |
| for unit_len in range(min_period, max_period + 1): | |
| if text[n - 1] != text[n - 1 - unit_len]: | |
| continue | |
| match_len = 1 | |
| idx = n - 2 | |
| while idx >= unit_len and text[idx] == text[idx - unit_len]: | |
| match_len += 1 | |
| idx -= 1 | |
| total_len = match_len + unit_len | |
| repeat_times = total_len // unit_len | |
| tail_len = total_len % unit_len | |
| if repeat_times >= min_repeat_times and total_len >= min_repeat_chars: | |
| return text[: n - total_len + unit_len] + text[n - tail_len:] | |
| return text | |
| def _filter_img_tags(text: str) -> str: | |
| """Drop placeholder `<img src="images/bbox_...">` blocks (model card).""" | |
| return "\n\n".join( | |
| block | |
| for block in text.split("\n\n") | |
| if not block.strip().startswith('<img src="images/bbox_') | |
| ) | |
| def _parse_region(region_text: str) -> str: | |
| """Normalize a user-supplied region into the card's compact JSON form. | |
| Accepts `[x1,y1,x2,y2]`, a JSON array, or whitespace-separated numbers. | |
| Args: | |
| region_text: Region input from the UI. | |
| Returns: | |
| A compact JSON array string like `"[100,100,900,300]"`. | |
| Raises: | |
| ValueError: If the input does not yield four integer coordinates. | |
| """ | |
| text = region_text.strip() | |
| if not text: | |
| raise ValueError("Region is empty. Provide four integers in [0, 1000).") | |
| try: | |
| values = json.loads(text) | |
| except json.JSONDecodeError: | |
| values = [tok for tok in re.split(r"[\s,]+", text) if tok] | |
| if not isinstance(values, list) or len(values) != 4: | |
| raise ValueError("Region must contain exactly four integers in [0, 1000).") | |
| region = [int(round(float(v))) for v in values] | |
| for v in region: | |
| if not 0 <= v <= 999: | |
| raise ValueError("Region coordinates must be integers in [0, 1000).") | |
| return json.dumps(region, separators=(",", ":")) | |
| # --------------------------------------------------------------------------- | |
| # Core inference — same contract as the card's `infer(...)`. | |
| # --------------------------------------------------------------------------- | |
| def _generate(images: list, prompt: str, max_new_tokens: int) -> str: | |
| """Run one greedy generation request over one or two images. | |
| Args: | |
| images: One or two RGB-converted PIL images (one request). | |
| prompt: Fully substituted task prompt. | |
| max_new_tokens: Generation token budget. | |
| Returns: | |
| The raw decoded completion text (unmodified). | |
| """ | |
| content = [{"type": "image", "image": img.convert("RGB")} for img in images] | |
| content.append({"type": "text", "text": prompt}) | |
| messages = [{"role": "user", "content": content}] | |
| inputs = processor.apply_chat_template( | |
| messages, | |
| tokenize=True, | |
| add_generation_prompt=True, | |
| return_dict=True, | |
| return_tensors="pt", | |
| enable_thinking=False, | |
| ).to(model.device) | |
| with torch.inference_mode(): | |
| generated = model.generate( | |
| **inputs, | |
| max_new_tokens=int(max_new_tokens), | |
| do_sample=False, | |
| temperature=None, | |
| top_p=None, | |
| top_k=None, | |
| ) | |
| trimmed = generated[:, inputs["input_ids"].shape[1]:] | |
| return processor.batch_decode( | |
| trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False | |
| )[0].strip() | |
| def _estimate_duration(*args, max_new_tokens=4096, **kwargs): | |
| """ZeroGPU duration estimator sized from live measurements (~30ms/token | |
| plus prefill overhead), covering the full requested token budget.""" | |
| return int(min(280, 15 + float(max_new_tokens) * 0.032)) | |
| # --------------------------------------------------------------------------- | |
| # Tab 1: single-image task engine (16 tasks across 4 domains). | |
| # --------------------------------------------------------------------------- | |
| TASKS = { | |
| # key: (label, prompt stem, task-input keys required) | |
| "document_markdown": ("Document → Markdown", "document_markdown", ()), | |
| "document_structured": ("Document → structured JSON", "document_structured", ()), | |
| "document_layout": ("Layout detection (JSON)", "document_layout", ()), | |
| "plain_ocr": ("Plain OCR", "plain_ocr", ()), | |
| "word_spotting": ("Word-level text spotting (JSON)", "word_spotting", ()), | |
| "reference_grounding": ("Reference grounding (JSON)", "reference_grounding", ("references",)), | |
| "region_ocr": ("Region OCR", "region_ocr", ("region",)), | |
| "formula_recognition": ("Formula recognition (LaTeX)", "formula_recognition", ()), | |
| "chart_parsing": ("Chart parsing (JSON)", "chart_parsing", ()), | |
| "text_translation": ("Text translation (region)", "text_translation", ("region", "language")), | |
| "image_classification": ("Image classification", "image_classification", ("target_description", "labels")), | |
| "visual_question_answering": ("Visual question answering", "visual_question_answering", ("question",)), | |
| "position_aware_vqa": ("Position-aware VQA (JSON)", "position_aware_vqa", ("question",)), | |
| "kie_open": ("KIE — open extraction", "kie_open", ()), | |
| "kie_fields": ("KIE — specified fields", "kie_fields", ("keys",)), | |
| "kie_schema": ("KIE — schema-based", "kie_schema", ("schema",)), | |
| "figure_captioning": ("Figure captioning (JSON)", "figure_captioning", ()), | |
| } | |
| # Task inputs shown per selected task; (input key, gr.Component) pairs. | |
| TASK_INPUT_KEYS = ( | |
| "question", | |
| "references", | |
| "region", | |
| "language", | |
| "labels", | |
| "target_description", | |
| "keys", | |
| "schema", | |
| ) | |
| DEFAULT_KEYS_JSON = '["invoice_number", "date", "total"]' | |
| DEFAULT_SCHEMA_JSON = json.dumps( | |
| { | |
| "invoice_number": "", | |
| "date": "", | |
| "seller": {"name": "", "address": ""}, | |
| "items": [{"name": "", "quantity": "", "unit_price": "", "amount": ""}], | |
| "total": "", | |
| }, | |
| ensure_ascii=False, | |
| ) | |
| DEFAULT_LABELS_JSON = '["financial report", "invoice", "letter", "academic paper"]' | |
| NO_TABLE_RESPONSE = '[{"source_pages": [], "html": "No cross-page table exists between these pages."}]' | |
| def _build_task_prompt( | |
| task: str, | |
| question: str, | |
| references: str, | |
| region: str, | |
| language: str, | |
| labels: str, | |
| target_description: str, | |
| keys: str, | |
| schema: str, | |
| ) -> str: | |
| """Build the fully substituted task prompt for a single-image task. | |
| Uses literal `.replace()` (never `str.format`) because prompts contain | |
| JSON and LaTeX braces. | |
| Args: | |
| task: Task key from TASKS. | |
| question: Question text for VQA tasks. | |
| references: JSON array of reference strings (grounding). | |
| region: Region spec `[x1,y1,x2,y2]` in [0,1000) coordinates. | |
| language: Target language for translation. | |
| labels: JSON array of candidate labels (classification). | |
| target_description: What to classify (classification). | |
| keys: JSON array of field names (KIE fields). | |
| schema: JSON output template (KIE schema). | |
| Returns: | |
| The fully substituted prompt string. | |
| """ | |
| if task not in TASKS: | |
| raise ValueError(f"Unknown task: {task}") | |
| if task == "visual_question_answering": | |
| if not question.strip(): | |
| raise ValueError("Provide a question for VQA tasks.") | |
| prompt = _general_ocr_prompt("visual_question_answering").replace("{question}", question.strip()) | |
| return re.sub(r"\n{3,}", "\n\n", prompt).strip() | |
| prompt = _general_ocr_prompt(task) if task in { | |
| "plain_ocr", "word_spotting", "reference_grounding", "region_ocr", | |
| "formula_recognition", "chart_parsing", "text_translation", | |
| "image_classification", "position_aware_vqa", | |
| } else load_prompt(task).strip() | |
| if task == "reference_grounding": | |
| refs = references.strip() or '["Invoice number", "Total"]' | |
| refs = json.dumps(json.loads(refs), ensure_ascii=False, separators=(",", ":")) | |
| prompt = prompt.replace("{references}", refs) | |
| elif task == "region_ocr": | |
| prompt = prompt.replace("{target_region}", _parse_region(region)) | |
| elif task == "text_translation": | |
| if not language.strip(): | |
| raise ValueError("Provide a target language for translation.") | |
| prompt = prompt.replace("{target_region}", _parse_region(region)) | |
| prompt = prompt.replace("{target_language}", language.strip()) | |
| elif task == "image_classification": | |
| if not target_description.strip(): | |
| raise ValueError("Provide a target description for classification.") | |
| labs = json.dumps( | |
| json.loads(labels) if labels.strip() else DEFAULT_LABELS_JSON, | |
| ensure_ascii=False, | |
| ) | |
| prompt = prompt.replace("{target description}", target_description.strip()) | |
| prompt = prompt.replace("{candidate_labels}", labs) | |
| elif task == "position_aware_vqa": | |
| if not question.strip(): | |
| raise ValueError("Provide a question for VQA tasks.") | |
| prompt = prompt.replace("{question}", question.strip()) | |
| elif task == "kie_fields": | |
| if not keys.strip(): | |
| raise ValueError("Provide a JSON array of field keys.") | |
| prompt = prompt.replace("{keys_json}", json.dumps(json.loads(keys), ensure_ascii=False)) | |
| elif task == "kie_schema": | |
| if not schema.strip(): | |
| raise ValueError("Provide a JSON output template.") | |
| prompt = prompt.replace("{schema_json}", json.dumps(json.loads(schema), ensure_ascii=False)) | |
| return prompt | |
| def run_task( | |
| image, | |
| task: str = "document_markdown", | |
| question: str = "What is the title of this document?", | |
| references: str = '["Invoice number", "Total"]', | |
| region: str = "[0,0,999,999]", | |
| language: str = "English", | |
| labels: str = DEFAULT_LABELS_JSON, | |
| target_description: str = "the document type", | |
| keys: str = DEFAULT_KEYS_JSON, | |
| schema: str = DEFAULT_SCHEMA_JSON, | |
| filter_image_tags: bool = True, | |
| max_new_tokens: int = 4096, | |
| ): | |
| """Run one of OvisOCR3's single-image document tasks on an uploaded image. | |
| Args: | |
| image: The document/page image to analyze. | |
| task: Which of the 16 supported tasks to run. | |
| question: Question text (VQA tasks). | |
| references: JSON array of references to ground (reference grounding). | |
| region: `[x1,y1,x2,y2]` region in [0,1000) coordinates (region OCR / translation). | |
| language: Target language (text translation). | |
| labels: JSON array of candidate labels (image classification). | |
| target_description: What to classify (image classification). | |
| keys: JSON array of field names (KIE specified fields). | |
| schema: JSON output template (KIE schema mode). | |
| filter_image_tags: Drop placeholder `<img src="images/bbox_...">` blocks (Markdown task only). | |
| max_new_tokens: Maximum number of tokens to generate. | |
| Returns: | |
| The model output: Markdown for document parsing tasks, JSON for | |
| structured tasks, plain text for OCR tasks. | |
| """ | |
| if image is None: | |
| raise ValueError("Upload a document image first.") | |
| max_new_tokens = max(256, min(int(max_new_tokens), MAX_NEW_TOKENS_CEILING)) | |
| prompt = _build_task_prompt( | |
| task, question, references, region, language, labels, | |
| target_description, keys, schema, | |
| ) | |
| raw = _generate([image], prompt, max_new_tokens) | |
| if task == "document_markdown": | |
| text = raw.strip() | |
| if filter_image_tags: | |
| text = _filter_img_tags(text) | |
| return _clean_truncated_repeats(text) | |
| return _clean_truncated_repeats(raw) | |
| def _run_task_ui( | |
| image, | |
| task: str = "document_markdown", | |
| question: str = "What is the title of this document?", | |
| references: str = '["Invoice number", "Total"]', | |
| region: str = "[0,0,999,999]", | |
| language: str = "English", | |
| labels: str = DEFAULT_LABELS_JSON, | |
| target_description: str = "the document type", | |
| keys: str = DEFAULT_KEYS_JSON, | |
| schema: str = DEFAULT_SCHEMA_JSON, | |
| filter_image_tags: bool = True, | |
| max_new_tokens: int = 4096, | |
| ): | |
| """UI wrapper for `run_task` that also fills the rendered Markdown preview.""" | |
| text = run_task( | |
| image, task, question, references, region, language, labels, | |
| target_description, keys, schema, filter_image_tags, max_new_tokens, | |
| ) | |
| return text, text | |
| # --------------------------------------------------------------------------- | |
| # Tab 2: cross-page table reconstruction (two-stage, per model card). | |
| # --------------------------------------------------------------------------- | |
| def _prepare_page_markdown(raw: str, index: int) -> str: | |
| """Turn a raw page-parse into the page-local Markdown used by Stage 2. | |
| Applies the card's pipeline: repeated-tail cleanup, code-fence stripping, | |
| and qualifying image references with the page prefix. | |
| Args: | |
| raw: Raw `document_markdown` generation for one page. | |
| index: 1-based page number. | |
| Returns: | |
| Page Markdown with page-qualified image references. | |
| """ | |
| markdown = _clean_truncated_repeats(raw.strip()) | |
| markdown = re.sub(r'^```(?:markdown|html)?\s*\n|\n```\s*$', '', markdown.strip(), flags=re.I) | |
| if not markdown: | |
| raise ValueError(f"Page {index}: empty Markdown") | |
| return re.sub( | |
| r'(?<![\w/])(?:page_\d+/)?(images/bbox_\d+_\d+_\d+_\d+\.jpg)', | |
| rf'page_{index}/\1', | |
| markdown, | |
| ) | |
| def parse_page(image, max_new_tokens: int = 4096): | |
| """Stage 1 helper: parse one page to Markdown (for cross-page reconstruction). | |
| Args: | |
| image: One page image. | |
| max_new_tokens: Maximum number of tokens to generate. | |
| Returns: | |
| Raw Markdown for the page (before Stage 2 post-processing). | |
| """ | |
| if image is None: | |
| raise ValueError("Upload a page image first.") | |
| max_new_tokens = max(256, min(int(max_new_tokens), MAX_NEW_TOKENS_CEILING)) | |
| return _generate([image], load_prompt("document_markdown").strip(), max_new_tokens) | |
| def _parse_page1(image, max_new_tokens): | |
| """UI handler: parse page 1 and store its Stage-2-ready Markdown.""" | |
| raw = parse_page(image, max_new_tokens) | |
| return _prepare_page_markdown(raw, 1) | |
| def _parse_page2(image, max_new_tokens): | |
| """UI handler: parse page 2 and store its Stage-2-ready Markdown.""" | |
| raw = parse_page(image, max_new_tokens) | |
| return _prepare_page_markdown(raw, 2) | |
| def reconstruct_crosspage_table( | |
| page1_image, | |
| page2_image, | |
| page1_markdown: str, | |
| page2_markdown: str, | |
| max_new_tokens: int = 4096, | |
| ): | |
| """Stage 2: reconstruct a table continued across two adjacent pages. | |
| Args: | |
| page1_image: The earlier page image. | |
| page2_image: The following page image. | |
| page1_markdown: Stage-1 Markdown for page 1 (from `parse_page`). | |
| page2_markdown: Stage-1 Markdown for page 2 (from `parse_page`). | |
| max_new_tokens: Maximum number of tokens to generate. | |
| Returns: | |
| A JSON array; each record has `source_pages` and an `html` table. | |
| When no cross-page table exists, the card's negative prediction is | |
| returned verbatim. | |
| """ | |
| max_new_tokens = max(256, min(int(max_new_tokens), MAX_NEW_TOKENS_CEILING)) | |
| pages = [img for img in (page1_image, page2_image) if img is not None] | |
| if len(pages) != 2: | |
| raise ValueError("Upload both page images for cross-page reconstruction.") | |
| if not page1_markdown.strip() or not page2_markdown.strip(): | |
| raise ValueError("Parse both pages first so their Markdown can be edited/reused.") | |
| md1 = _prepare_page_markdown(page1_markdown, 1) | |
| md2 = _prepare_page_markdown(page2_markdown, 2) | |
| has_table = re.compile(r'<table\b[^>]*>.*?</table\s*>', re.I | re.S) | |
| if not (has_table.search(md1) or has_table.search(md2)): | |
| return NO_TABLE_RESPONSE | |
| page_context = "\n\n".join( | |
| f'<INPUT_PAGE index="{index}">\n{markdown}\n</INPUT_PAGE>' | |
| for index, markdown in enumerate((md1, md2), start=1) | |
| ) | |
| prompt = load_prompt("crosspage_table").removesuffix("\n") + "\n\nInput Markdown:\n\n" + page_context | |
| raw = _generate(pages, prompt, max_new_tokens) | |
| return _clean_truncated_repeats(raw) | |
| def _reconstruct_ui(page1_image, page2_image, page1_markdown, page2_markdown, max_new_tokens): | |
| """UI wrapper for cross-page reconstruction; returns the JSON result text.""" | |
| return reconstruct_crosspage_table( | |
| page1_image, page2_image, page1_markdown, page2_markdown, max_new_tokens | |
| ) | |
| # --------------------------------------------------------------------------- | |
| # UI. | |
| # --------------------------------------------------------------------------- | |
| CSS = """ | |
| #col-container { max-width: 1280px; margin: 0 auto; } | |
| .dark .gradio-container { color: var(--body-text-color); } | |
| """ | |
| TASK_TO_INPUTS = {key: spec[2] for key, spec in TASKS.items()} | |
| def _toggle_task_inputs(task: str): | |
| """Return per-task-input visibility flags for the selected task. | |
| Args: | |
| task: Task key from the dropdown. | |
| Returns: | |
| A list of `gr.update` visibility values, ordered as TASK_INPUT_KEYS. | |
| """ | |
| required = set(TASK_TO_INPUTS.get(task, ())) | |
| return [ | |
| gr.update(visible=key in required) for key in TASK_INPUT_KEYS | |
| ] | |
| with gr.Blocks() as demo: | |
| with gr.Column(elem_id="col-container"): | |
| gr.Markdown( | |
| """ | |
| # 📑 OvisOCR3 — Unified Document Understanding | |
| Interactive demo for [**ATH-MaaS/OvisOCR3**](https://huggingface.co/ATH-MaaS/OvisOCR3), | |
| a compact 0.8B vision-language model that covers **document parsing**, **general | |
| OCR**, **key information extraction**, **figure captioning**, and **cross-page | |
| table reconstruction** — all selected through task prompts on the same checkpoint. | |
| Upload a page and pick a task. Task inputs appear only where the task needs them. | |
| """ | |
| ) | |
| with gr.Tabs(): | |
| # ------------------------- Tab 1: tasks ------------------------- | |
| with gr.TabItem("Document tasks"): | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| image_in = gr.Image(type="pil", label="Document image", height=460) | |
| task_dd = gr.Dropdown( | |
| choices=[(label, key) for key, (label, _, _) in TASKS.items()], | |
| value="document_markdown", | |
| label="Task", | |
| filterable=False, | |
| ) | |
| run_btn = gr.Button("Run task", variant="primary") | |
| with gr.Accordion("Task inputs", open=True): | |
| question_tb = gr.Textbox( | |
| label="Question (VQA tasks)", | |
| value="What is the title of this document?", | |
| lines=2, | |
| ) | |
| references_tb = gr.Textbox( | |
| label="References — JSON array (reference grounding)", | |
| value='["Invoice number", "Total"]', | |
| lines=2, | |
| ) | |
| region_tb = gr.Textbox( | |
| label="Region — [x1,y1,x2,y2] in [0,1000) (region OCR / translation)", | |
| value="[0,0,999,999]", | |
| lines=1, | |
| ) | |
| language_tb = gr.Textbox( | |
| label="Target language (translation)", | |
| value="English", | |
| lines=1, | |
| ) | |
| labels_tb = gr.Textbox( | |
| label="Candidate labels — JSON array (classification)", | |
| value=DEFAULT_LABELS_JSON, | |
| lines=2, | |
| ) | |
| target_desc_tb = gr.Textbox( | |
| label="Target description (classification)", | |
| value="the document type", | |
| lines=1, | |
| ) | |
| keys_tb = gr.Textbox( | |
| label=f"Field keys — JSON array (KIE fields). Example: {DEFAULT_KEYS_JSON}", | |
| value=DEFAULT_KEYS_JSON, | |
| lines=2, | |
| ) | |
| schema_tb = gr.Textbox( | |
| label="Output template — JSON (KIE schema). Empty strings are value placeholders.", | |
| value=DEFAULT_SCHEMA_JSON, | |
| lines=6, | |
| ) | |
| with gr.Accordion("Advanced settings", open=False): | |
| filter_tags = gr.Checkbox( | |
| value=True, | |
| label="Filter figure placeholder <img> tags (Markdown task only)", | |
| ) | |
| max_tokens = gr.Slider( | |
| 512, MAX_NEW_TOKENS_CEILING, value=4096, step=512, | |
| label="Max new tokens", | |
| ) | |
| with gr.Column(scale=1): | |
| task_out = gr.Code(label="Output", lines=24) | |
| with gr.Accordion("Rendered Markdown preview", open=False): | |
| task_render = gr.Markdown() | |
| task_dd.change( | |
| _toggle_task_inputs, | |
| inputs=[task_dd], | |
| outputs=[ | |
| question_tb, references_tb, region_tb, language_tb, | |
| labels_tb, target_desc_tb, keys_tb, schema_tb, | |
| ], | |
| ) | |
| # Apply visibility for the default task on load. | |
| demo.load( | |
| _toggle_task_inputs, | |
| inputs=[task_dd], | |
| outputs=[ | |
| question_tb, references_tb, region_tb, language_tb, | |
| labels_tb, target_desc_tb, keys_tb, schema_tb, | |
| ], | |
| ) | |
| task_inputs = [ | |
| image_in, task_dd, question_tb, references_tb, region_tb, | |
| language_tb, labels_tb, target_desc_tb, keys_tb, schema_tb, | |
| filter_tags, max_tokens, | |
| ] | |
| run_btn.click( | |
| _run_task_ui, | |
| inputs=task_inputs, | |
| outputs=[task_out, task_render], | |
| api_name="run_task", | |
| ) | |
| gr.Examples( | |
| examples=[ | |
| ["examples/paper_page.png", "document_markdown"], | |
| ["examples/financial_table.jpg", "kie_open"], | |
| ["examples/handwriting.jpg", "plain_ocr"], | |
| ["examples/math_textbook.png", "formula_recognition"], | |
| ["examples/magazine_page.png", "figure_captioning"], | |
| ["examples/paper_page.png", "document_structured"], | |
| ], | |
| inputs=[image_in, task_dd], | |
| outputs=[task_out, task_render], | |
| fn=_run_task_ui, | |
| cache_examples=True, | |
| cache_mode="lazy", | |
| ) | |
| # --------------------- Tab 2: cross-page tables --------------------- | |
| with gr.TabItem("Cross-page table"): | |
| gr.Markdown( | |
| """ | |
| **Two-stage cross-page table reconstruction** (one request per stage): | |
| 1. Parse each page to Markdown (Stage 1 — buttons below, one call per page). | |
| 2. Reconstruct the table that continues across the page boundary (Stage 2 — | |
| one two-image request with both pages and their Markdown). | |
| Both page images are required. You can edit the parsed Markdown before | |
| Stage 2. | |
| """ | |
| ) | |
| with gr.Row(): | |
| with gr.Column(): | |
| cp_page1_img = gr.Image(type="pil", label="Page 1 (earlier page)", height=320) | |
| cp_parse1_btn = gr.Button("Parse page 1 → Markdown", variant="secondary") | |
| cp_page1_md = gr.Code(label="Page 1 Markdown", language="markdown", lines=10) | |
| with gr.Column(): | |
| cp_page2_img = gr.Image(type="pil", label="Page 2 (following page)", height=320) | |
| cp_parse2_btn = gr.Button("Parse page 2 → Markdown", variant="secondary") | |
| cp_page2_md = gr.Code(label="Page 2 Markdown", language="markdown", lines=10) | |
| with gr.Accordion("Advanced settings", open=False): | |
| cp_max_tokens = gr.Slider( | |
| 512, MAX_NEW_TOKENS_CEILING, value=4096, step=512, | |
| label="Max new tokens (Stage 1 parsing)", | |
| ) | |
| cp_max_tokens2 = gr.Slider( | |
| 512, MAX_NEW_TOKENS_CEILING, value=4096, step=512, | |
| label="Max new tokens (Stage 2 reconstruction)", | |
| ) | |
| cp_run_btn = gr.Button("Reconstruct cross-page table", variant="primary") | |
| cp_out = gr.Code(label="Reconstructed table (JSON)", language="json", lines=12) | |
| cp_parse1_btn.click( | |
| _parse_page1, | |
| inputs=[cp_page1_img, cp_max_tokens], | |
| outputs=[cp_page1_md], | |
| api_name="parse_page1", | |
| ) | |
| cp_parse2_btn.click( | |
| _parse_page2, | |
| inputs=[cp_page2_img, cp_max_tokens], | |
| outputs=[cp_page2_md], | |
| api_name="parse_page2", | |
| ) | |
| cp_run_btn.click( | |
| _reconstruct_ui, | |
| inputs=[cp_page1_img, cp_page2_img, cp_page1_md, cp_page2_md, cp_max_tokens2], | |
| outputs=[cp_out], | |
| api_name="reconstruct_crosspage_table", | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True) |