adelevett/docling_pp_layout_demo

1# ---------------------------------------------------------------------------2# Force-upgrade transformers to >=5.1.0 before any other import.3#4# Why: PP-DocLayoutV3's custom model classes (PPDocLayoutV3ImageProcessor,5# PPDocLayoutV3ForObjectDetection) were added to the transformers library in6# version 5.1.0.  docling-ibm-models caps transformers<5.0.0 (conservative7# pinning), so pip resolves transformers ~4.x at build time.  We upgrade it8# here at runtime, before any docling/transformers import, so the correct9# classes are available.  docling-ibm-models' usage (AutoModel, pipeline API)10# remains compatible with transformers 5.x.11# ---------------------------------------------------------------------------12import subprocess13import sys14 15subprocess.run(16    [17        sys.executable, "-m", "pip", "install",18        "transformers>=5.1.0",19        "--quiet",20    ],21    check=True,22)23 24# `spaces` MUST be imported before any package that touches CUDA (torch,25# transformers, docling …).  ZeroGPU intercepts the CUDA initialisation; if26# anything else triggers it first the import raises RuntimeError.27import spaces  # noqa: E40228 29# ---------------------------------------------------------------------------30# Plugin registration31# ---------------------------------------------------------------------------32# docling-pp-doc-layout requires Python >=3.12 on PyPI, but the code itself33# is compatible with Python 3.10 (all annotations are guarded by34# `from __future__ import annotations`).  Instead of installing the package,35# we bundle the source directly and register the model with docling's factory36# by monkey-patching BaseFactory.load_from_plugins so that every new37# LayoutFactory instance automatically includes PPDocLayoutV3Model.38from docling.models.factories.base_factory import BaseFactory39from docling.models.factories.layout_factory import LayoutFactory40from docling_pp_doc_layout.model import PPDocLayoutV3Model41 42_orig_load = BaseFactory.load_from_plugins43 44 45def _load_with_pp_doc_layout(46    self, plugin_name=None, allow_external_plugins=False47):48    _orig_load(49        self,50        plugin_name=plugin_name,51        allow_external_plugins=allow_external_plugins,52    )53    if isinstance(self, LayoutFactory):54        try:55            self.register(56                PPDocLayoutV3Model,57                "docling-pp-doc-layout",58                "docling_pp_doc_layout.model",59            )60        except ValueError:61            pass  # already registered on a previous factory creation62 63 64BaseFactory.load_from_plugins = _load_with_pp_doc_layout65 66# ---------------------------------------------------------------------------67import gradio as gr68from docling.datamodel.base_models import InputFormat69from docling.document_converter import DocumentConverter, PdfFormatOption70from docling.datamodel.pipeline_options import PdfPipelineOptions71from docling_pp_doc_layout.options import PPDocLayoutV3Options72 73# Global initialisation — pipeline is constructed lazily on the first74# convert() call, which happens inside @spaces.GPU, so decide_device()75# correctly resolves "cuda:0" when the H200 is allocated.76pipeline_options = PdfPipelineOptions(77    layout_options=PPDocLayoutV3Options(78        batch_size=2,79        confidence_threshold=0.5,80    )81)82 83converter = DocumentConverter(84    format_options={85        InputFormat.PDF: PdfFormatOption(pipeline_options=pipeline_options)86    }87)88 89 90@spaces.GPU(duration=120)91def infer_layout(file_path: str | None):92    if not file_path:93        return {"error": "No file uploaded"}, None94    try:95        result = converter.convert(file_path)96        structured_data = []97        for item, _level in result.document.iterate_items():98            structured_data.append({99                "type": type(item).__name__,100                "content": getattr(item, "text", "No text mapping"),101            })102        # Write to a temp file so Gradio can serve it as a download.103        import json, tempfile, os104        tmp = tempfile.NamedTemporaryFile(105            mode="w", suffix=".json", delete=False, encoding="utf-8"106        )107        json.dump(structured_data, tmp, ensure_ascii=False, indent=2)108        tmp.close()109        return structured_data, tmp.name110    except Exception as e:111        return {"runtime_exception": str(e)}, None112 113 114with gr.Blocks(title="PP-DocLayoutV3 Empirical Parser") as interface:115    gr.Markdown(116        "## Layout Detection Inference\n"117        "Upload a PDF to parse structural components through the "118        "PaddlePaddle PP-DocLayoutV3 model."119    )120    with gr.Row():121        pdf_input = gr.File(label="Source Document", file_types=[".pdf"])122        json_output = gr.JSON(label="Structured Extraction Matrix")123    download_btn = gr.DownloadButton(label="Download JSON", visible=False)124    execute_btn = gr.Button("Run Layout Detection")125 126    def run_and_reveal(file_path):127        data, path = infer_layout(file_path)128        return data, gr.DownloadButton(value=path, visible=path is not None)129 130    execute_btn.click(131        fn=run_and_reveal,132        inputs=pdf_input,133        outputs=[json_output, download_btn],134    )135 136if __name__ == "__main__":137    interface.launch()