parse-bench 1.0.0__py3-none-any.whl

parse_bench/inference/providers/layoutdet/yolo.py

@@ -0,0 +1,109 @@
+"""Provider for YOLO-DocLayNet layout detection."""
+
+from typing import Any
+
+from parse_bench.inference.providers.base import ProviderPermanentError
+from parse_bench.inference.providers.layoutdet.base import HFLayoutDetProvider
+from parse_bench.inference.providers.registry import register_provider
+from parse_bench.schemas.layout_detection_output import (
+    LayoutDetectionModel,
+    LayoutOutput,
+    LayoutPrediction,
+    YoloLabel,
+)
+from parse_bench.schemas.pipeline_io import InferenceResult, RawInferenceResult
+from parse_bench.schemas.product import ProductType
+
+
+@register_provider("yolo_layout")
+class YoloLayoutProvider(HFLayoutDetProvider):
+    """
+    Provider for YOLO-DocLayNet layout detection model.
+
+    This provider uses the YOLO model trained on DocLayNet served on HuggingFace
+    inference endpoints for detecting document layout regions.
+
+    Response format:
+        {
+            "pred_boxes": [[x1, y1, x2, y2], ...],
+            "pred_classes": [class_id, ...],
+            "scores": [score, ...]
+        }
+    """
+
+    endpoint_url = "https://exqoktya7l52qu1r.us-east-1.aws.endpoints.huggingface.cloud"
+    model_type = LayoutDetectionModel.YOLO_DOCLAYNET
+
+    def __init__(
+        self,
+        provider_name: str,
+        base_config: dict[str, Any] | None = None,
+    ):
+        """Initialize the YOLO-DocLayNet layout detection provider."""
+        super().__init__(provider_name, base_config)
+
+    def _parse_response(self, response: dict[str, Any]) -> list[LayoutPrediction]:
+        """
+        Parse YOLO-DocLayNet response into layout predictions.
+
+        :param response: Raw JSON response with pred_boxes, pred_classes, scores
+        :return: List of unified LayoutPrediction objects
+        """
+        predictions: list[LayoutPrediction] = []
+
+        boxes = response.get("pred_boxes", [])
+        classes = response.get("pred_classes", [])
+        scores = response.get("scores", [])
+
+        for bbox, class_id, score in zip(boxes, classes, scores, strict=False):
+            # Convert class_id to YoloLabel enum
+            # Model outputs 0-indexed labels (0-10) that match YoloLabel directly
+            label = YoloLabel(class_id)
+            predictions.append(
+                LayoutPrediction(
+                    bbox=bbox,
+                    score=score,
+                    label=str(int(label)),
+                    provider_metadata={"label_name": label.name},
+                )
+            )
+
+        return predictions
+
+    def normalize(self, raw_result: RawInferenceResult) -> InferenceResult:
+        """
+        Normalize raw inference result to produce LayoutOutput.
+
+        :param raw_result: Raw inference result from run_inference()
+        :return: Inference result with both raw and normalized outputs
+        :raises ProviderError: For any normalization failures
+        """
+        if raw_result.product_type != ProductType.LAYOUT_DETECTION:
+            raise ProviderPermanentError(
+                f"{self.__class__.__name__} only supports LAYOUT_DETECTION product type, got {raw_result.product_type}"
+            )
+
+        # Parse the response into raw predictions
+        response = raw_result.raw_output.get("response", {})
+        raw_predictions = self._parse_response(response)
+
+        output = LayoutOutput(
+            task_type="layout_detection",
+            example_id=raw_result.request.example_id,
+            pipeline_name=raw_result.pipeline_name,
+            model=self.model_type,
+            image_width=max(int(raw_result.raw_output.get("image_width", 1)), 1),
+            image_height=max(int(raw_result.raw_output.get("image_height", 1)), 1),
+            predictions=raw_predictions,
+        )
+
+        return InferenceResult(
+            request=raw_result.request,
+            pipeline_name=raw_result.pipeline_name,
+            product_type=raw_result.product_type,
+            raw_output=raw_result.raw_output,
+            output=output,
+            started_at=raw_result.started_at,
+            completed_at=raw_result.completed_at,
+            latency_in_ms=raw_result.latency_in_ms,
+        )

parse_bench/inference/providers/parse/__init__.py

@@ -0,0 +1,64 @@
+"""Parse providers — imported lazily to avoid requiring all SDKs."""
+
+import importlib
+import logging
+
+logger = logging.getLogger(__name__)
+
+_PROVIDER_MODULES = [
+    "amazon_nova",
+    "anthropic",
+    "azure_document_intelligence",
+    "chandra2",
+    "chunkr",
+    "databricks_ai_parse",
+    "datalab",
+    "deepseekocr2",
+    "docling",
+    "docling_serve",
+    "dots_ocr",
+    "extend_parse",
+    "falconocr",
+    "florin_parser_nano",
+    "gemma4",
+    "glm_zai",
+    "google",
+    "google_docai",
+    "granite_vision",
+    "infinity_parser2",
+    "kdl_frontier_nano",
+    "landingai",
+    "liteparse",
+    "markitdown",
+    "opendataloader",
+    "pdf_inspector",
+    "pymupdf4llm",
+    "rakedoc_nano",
+    "llamaparse",
+    "llamaparse_v2_normalization",
+    "mineru25",
+    "mineru2605pro",
+    "mineru_diffusion",
+    "mistral_ocr",
+    "nemotron_omni",
+    "openai",
+    "paddleocr",
+    "pulse",
+    "pymupdf",
+    "pypdf",
+    "qwen",
+    "reducto",
+    "surya2",
+    "tesseract",
+    "textract",
+    "unlimitedocr",
+    "unstructured",
+    "warp_ingest",
+    "oi_parser",
+]
+
+for _mod in _PROVIDER_MODULES:
+    try:
+        importlib.import_module(f"parse_bench.inference.providers.parse.{_mod}")
+    except ImportError:
+        logger.debug("Skipping parse provider %s (missing dependency)", _mod)

parse_bench/inference/providers/parse/_docling_common.py

@@ -0,0 +1,233 @@
+"""Common functionality for docling and docling_serve providers."""
+
+from typing import Any
+
+from docling_core.types.doc.document import DoclingDocument
+
+from parse_bench.layout_label_mapping import (
+    UnknownRawLayoutLabelError,
+    map_docling_raw_label_to_canonical,
+)
+from parse_bench.schemas.parse_output import (
+    LayoutItemIR,
+    LayoutSegmentIR,
+    ParseLayoutPageIR,
+)
+
+_DOCLING_EXCLUDED_LAYOUT_LABELS = frozenset(
+    {
+        "empty_value",
+        "field_heading",
+        "field_hint",
+        "field_item",
+        "field_key",
+        "field_region",
+        "field_value",
+        "marker",
+    }
+)
+_DOCLING_TABLE_LABELS = frozenset({"document_index", "table"})
+_DOCLING_IMAGE_LABELS = frozenset({"chart", "picture"})
+
+
+def _normalize_docling_label(label: object) -> str | None:
+    if label is None:
+        return None
+    value = getattr(label, "value", label)
+    if not isinstance(value, str):
+        return None
+    return value.strip().lower()
+
+
+def _should_include_docling_label(raw_label: str) -> bool:
+    if raw_label in _DOCLING_EXCLUDED_LAYOUT_LABELS:
+        return False
+    try:
+        map_docling_raw_label_to_canonical(raw_label)
+    except UnknownRawLayoutLabelError:
+        return False
+    return True
+
+
+def _docling_item_type(raw_label: str) -> str:
+    if raw_label in _DOCLING_TABLE_LABELS:
+        return "table"
+    if raw_label in _DOCLING_IMAGE_LABELS:
+        return "image"
+    return "text"
+
+
+def _extract_docling_item_value(item: Any, doc: DoclingDocument, raw_label: str) -> str:
+    item_type = _docling_item_type(raw_label)
+    if item_type == "image":
+        return ""
+
+    if item_type == "table" and hasattr(item, "export_to_html"):
+        try:
+            html = item.export_to_html(doc=doc, add_caption=True)
+            if isinstance(html, str):
+                return html
+        except Exception:
+            pass
+
+    text = getattr(item, "text", None)
+    if isinstance(text, str):
+        return text
+
+    if hasattr(item, "export_to_markdown"):
+        try:
+            markdown = item.export_to_markdown()
+            if isinstance(markdown, str):
+                return markdown
+        except Exception:
+            pass
+
+    return ""
+
+
+def _normalize_docling_charspan(
+    charspan: object,
+    *,
+    text_length: int,
+    include_span: bool,
+) -> tuple[int | None, int | None]:
+    if not include_span or not isinstance(charspan, (list, tuple)) or len(charspan) != 2:
+        return (None, None)
+
+    start_raw, end_raw = charspan
+    if not isinstance(start_raw, int) or not isinstance(end_raw, int):
+        return (None, None)
+
+    start = max(0, min(start_raw, text_length))
+    end_exclusive = max(start, min(end_raw, text_length))
+    if end_exclusive <= start:
+        return (None, None)
+
+    # Docling charspan behaves like a Python slice [start, end).
+    return (start, end_exclusive - 1)
+
+
+def _build_docling_segment(
+    *,
+    prov: Any,
+    raw_label: str,
+    page_width: float,
+    page_height: float,
+    include_span: bool,
+    text_length: int,
+) -> LayoutSegmentIR | None:
+    bbox = getattr(prov, "bbox", None)
+    if bbox is None or page_width <= 0 or page_height <= 0:
+        return None
+
+    bbox_top_left = bbox.to_top_left_origin(page_height=page_height)
+    width = bbox_top_left.r - bbox_top_left.l
+    height = bbox_top_left.b - bbox_top_left.t
+    if width <= 0 or height <= 0:
+        return None
+
+    start_index, end_index = _normalize_docling_charspan(
+        getattr(prov, "charspan", None),
+        text_length=text_length,
+        include_span=include_span,
+    )
+
+    return LayoutSegmentIR(
+        x=bbox_top_left.l / page_width,
+        y=bbox_top_left.t / page_height,
+        w=width / page_width,
+        h=height / page_height,
+        confidence=1.0,
+        label=raw_label,
+        start_index=start_index,
+        end_index=end_index,
+    )
+
+
+def _merge_segments(segments: list[LayoutSegmentIR]) -> LayoutSegmentIR | None:
+    if not segments:
+        return None
+
+    x1 = min(segment.x for segment in segments)
+    y1 = min(segment.y for segment in segments)
+    x2 = max(segment.x + segment.w for segment in segments)
+    y2 = max(segment.y + segment.h for segment in segments)
+    return LayoutSegmentIR(
+        x=x1,
+        y=y1,
+        w=x2 - x1,
+        h=y2 - y1,
+        confidence=1.0,
+        label=segments[0].label,
+    )
+
+
+def _build_docling_layout_pages(
+    *,
+    doc: DoclingDocument,
+    raw_pages: list[dict[str, Any]],
+) -> list[ParseLayoutPageIR]:
+    page_markdown_by_number: dict[int, str] = {}
+    for page_data in raw_pages:
+        page_number = page_data.get("page")
+        if isinstance(page_number, int) and page_number > 0:
+            page_markdown_by_number[page_number] = str(page_data.get("markdown", ""))
+
+    layout_pages: list[ParseLayoutPageIR] = []
+    for page_number in sorted(doc.pages.keys()):
+        page = doc.pages[page_number]
+        page_width = float(page.size.width)
+        page_height = float(page.size.height)
+        items: list[LayoutItemIR] = []
+
+        for item, _level in doc.iterate_items(page_no=page_number):
+            raw_label = _normalize_docling_label(getattr(item, "label", None))
+            if raw_label is None or not _should_include_docling_label(raw_label):
+                continue
+
+            item_type = _docling_item_type(raw_label)
+            item_value = _extract_docling_item_value(item, doc, raw_label)
+            include_span = item_type == "text"
+
+            page_provs = [
+                prov for prov in getattr(item, "prov", []) or [] if getattr(prov, "page_no", None) == page_number
+            ]
+            segments = [
+                segment
+                for prov in page_provs
+                if (
+                    segment := _build_docling_segment(
+                        prov=prov,
+                        raw_label=raw_label,
+                        page_width=page_width,
+                        page_height=page_height,
+                        include_span=include_span,
+                        text_length=len(item_value),
+                    )
+                )
+                is not None
+            ]
+            if not segments:
+                continue
+
+            merged_bbox = _merge_segments(segments)
+            items.append(
+                LayoutItemIR(
+                    type=item_type,
+                    value=item_value,
+                    bbox=merged_bbox,
+                    layout_segments=segments,
+                )
+            )
+
+        layout_pages.append(
+            ParseLayoutPageIR(
+                page_number=page_number,
+                width=page_width,
+                height=page_height,
+                md=page_markdown_by_number.get(page_number, ""),
+                items=items,
+            )
+        )
+
+    return layout_pages