Spaces:
Running on Zero
Running on Zero
| """Docling document parsing with figure extraction and markdown export.""" | |
| from __future__ import annotations | |
| import os | |
| import uuid | |
| from collections.abc import Callable | |
| from pathlib import Path | |
| from typing import Any | |
| from storage import get_temp_dir | |
| def _top_class(figure: Any) -> str | None: | |
| """Return the highest-confidence Docling classification label for a picture item. | |
| Tries ``figure.meta.classification.predictions`` first (docling-core ≥2.5), | |
| then falls back to the deprecated ``figure.annotations`` list. | |
| Returns ``None`` if classification is unavailable or fails. | |
| """ | |
| try: | |
| meta = getattr(figure, "meta", None) | |
| if meta is not None: | |
| classification = getattr(meta, "classification", None) | |
| if classification is not None: | |
| preds = getattr(classification, "predictions", None) | |
| if preds: | |
| best = max(preds, key=lambda p: p.confidence) | |
| return best.class_name | |
| except Exception: # noqa: BLE001 | |
| pass | |
| try: | |
| for ann in getattr(figure, "annotations", []): | |
| preds = getattr(ann, "predicted_classes", None) | |
| if preds: | |
| best = max(preds, key=lambda p: p.confidence) | |
| return best.class_name | |
| except Exception: # noqa: BLE001 | |
| pass | |
| return None | |
| _EXT_TO_INPUT_FORMAT = { | |
| ".pdf": "PDF", | |
| ".docx": "DOCX", | |
| ".xlsx": "XLSX", | |
| ".pptx": "PPTX", | |
| } | |
| _converter: Any = None | |
| def _is_zerogpu_space() -> bool: | |
| """Detect HuggingFace ZeroGPU Spaces where CUDA init must happen inside @spaces.GPU.""" | |
| return bool(os.environ.get("SPACE_ID")) | |
| def _pick_accelerator_device() -> Any: | |
| """Return the best AcceleratorDevice for Docling, or None if unsupported.""" | |
| try: | |
| from docling.datamodel.pipeline_options import AcceleratorDevice | |
| except ImportError: | |
| return None | |
| if _is_zerogpu_space(): | |
| return AcceleratorDevice.CPU | |
| try: | |
| import torch | |
| if torch.cuda.is_available(): | |
| return AcceleratorDevice.CUDA | |
| if torch.backends.mps.is_available(): | |
| return AcceleratorDevice.MPS | |
| except ImportError: | |
| pass | |
| return AcceleratorDevice.CPU | |
| def get_converter() -> Any: | |
| """Return a shared DocumentConverter instance, creating it on first call.""" | |
| global _converter # noqa: PLW0603 | |
| if _converter is not None: | |
| return _converter | |
| from docling.datamodel.base_models import InputFormat | |
| from docling.document_converter import DocumentConverter, PdfFormatOption | |
| pdf_format_option = PdfFormatOption() | |
| pdf_format_option.pipeline_options.generate_picture_images = True | |
| pdf_format_option.pipeline_options.images_scale = 2.0 | |
| pdf_format_option.pipeline_options.do_picture_classification = True | |
| device = _pick_accelerator_device() | |
| if device is not None: | |
| try: | |
| from docling.datamodel.pipeline_options import AcceleratorOptions | |
| pdf_format_option.pipeline_options.accelerator_options = AcceleratorOptions( | |
| device=device, | |
| ) | |
| except Exception: # noqa: BLE001 | |
| pass | |
| # On ZeroGPU, hide CUDA from sub-libraries (RapidOCR, etc.) that do their | |
| # own torch.cuda detection and trigger real CUDA init. | |
| hide_cuda = _is_zerogpu_space() | |
| old_val = os.environ.get("CUDA_VISIBLE_DEVICES") | |
| if hide_cuda: | |
| os.environ["CUDA_VISIBLE_DEVICES"] = "" | |
| try: | |
| _converter = DocumentConverter(format_options={InputFormat.PDF: pdf_format_option}) | |
| finally: | |
| if hide_cuda: | |
| if old_val is None: | |
| os.environ.pop("CUDA_VISIBLE_DEVICES", None) | |
| else: | |
| os.environ["CUDA_VISIBLE_DEVICES"] = old_val | |
| return _converter | |
| def parse_document( | |
| file_bytes: bytes, | |
| file_ext: str = ".pdf", | |
| on_progress: Callable[[str], None] | None = None, | |
| ) -> dict[str, Any]: | |
| """Parse a document with Docling and extract markdown, text, and figure regions. | |
| Args: | |
| file_bytes: Document file content as bytes. | |
| file_ext: File extension (e.g. ``".pdf"``, ``".docx"``, ``".xlsx"``, ``".pptx"``). | |
| on_progress: Optional callback ``(phase_message) -> None`` for progress reporting. | |
| Returns: | |
| Dictionary with keys: | |
| - ``html``: HTML-wrapped markdown representation of the document. | |
| - ``text``: Full extracted plain text. | |
| - ``figures``: List of figure dicts with ``bbox``, ``page``, ``caption``, ``image``, | |
| and ``figure_type`` (top Docling classification label or ``None``). | |
| """ | |
| def _report(msg: str) -> None: | |
| if on_progress: | |
| on_progress(msg) | |
| try: | |
| from docling.datamodel.base_models import InputFormat | |
| ext = file_ext.lower() | |
| tmp_path = str(get_temp_dir() / f"{uuid.uuid4().hex}{ext}") | |
| Path(tmp_path).write_bytes(file_bytes) | |
| try: | |
| format_name = _EXT_TO_INPUT_FORMAT.get(ext, "PDF") | |
| input_format = getattr(InputFormat, format_name) | |
| _report("Initializing document converter...") | |
| converter = get_converter() | |
| _report("Converting document (this may take a moment)...") | |
| result = converter.convert(tmp_path) | |
| doc = result.document | |
| _report("Exporting document content...") | |
| markdown_text = doc.export_to_markdown() | |
| html = markdown_text | |
| text = doc.export_to_text() | |
| _report("Processing figures...") | |
| figures: list[dict[str, Any]] = [] | |
| try: | |
| if hasattr(doc, "pictures"): | |
| for figure in doc.pictures: | |
| if figure.content_layer.value != "body": | |
| continue | |
| page_num = 0 | |
| bbox_list = None | |
| if figure.prov: | |
| page_num = figure.prov[0].page_no - 1 # Docling is 1-based | |
| bbox = figure.prov[0].bbox | |
| bbox_list = [bbox.l, bbox.t, bbox.width, bbox.height] | |
| caption = "" | |
| if figure.captions: | |
| for cap_ref in figure.captions: | |
| try: | |
| if hasattr(cap_ref, "cref") and cap_ref.cref.startswith("#/texts/"): | |
| idx = int(cap_ref.cref.split("/")[-1]) | |
| if idx < len(doc.texts): | |
| caption = doc.texts[idx].text | |
| break | |
| except Exception: # noqa: BLE001 | |
| pass | |
| if figure.image: | |
| try: | |
| pil_image = figure.image.pil_image | |
| figures.append({ | |
| "bbox": bbox_list, | |
| "page": page_num, | |
| "caption": caption, | |
| "image": pil_image, | |
| "figure_type": _top_class(figure), | |
| }) | |
| except Exception: # noqa: BLE001 | |
| pass | |
| except Exception: # noqa: BLE001 | |
| figures = [] | |
| return {"html": html, "text": text, "figures": figures} | |
| finally: | |
| if os.path.exists(tmp_path): | |
| os.unlink(tmp_path) | |
| except ImportError as e: | |
| print(f"Docling import error: {e}, using placeholder") | |
| return { | |
| "html": "<h1>Sample Document</h1><p>Docling not available - using placeholder.</p>", | |
| "text": "Sample text from PDF.\n\nDocling not available - using placeholder.", | |
| "figures": [], | |
| } | |
| except Exception as e: # noqa: BLE001 | |
| import traceback | |
| print(f"Docling parse error: {e}") | |
| traceback.print_exc() | |
| return { | |
| "html": f"<h1>Error</h1><pre>{e!s}</pre>", | |
| "text": f"Error: {e!s}", | |
| "figures": [], | |
| } | |