parse-bench 1.0.0__py3-none-any.whl

parse_bench/inference/providers/parse/deepseekocr2.py

@@ -0,0 +1,382 @@
+"""Provider for DeepSeek-OCR-2 Modal server.
+
+DeepSeek-OCR-2 (deepseek-ai/DeepSeek-OCR-2) is a MoE vision-language model
+that handles layout detection + OCR in a single pass via the <|grounding|> token.
+
+API format: POST /predict with {"image_base64": "..."} → {"markdown": "...", "status": "success"}
+"""
+
+import asyncio
+import base64
+import io
+import os
+import re
+from datetime import datetime
+from pathlib import Path
+from typing import Any
+
+import aiohttp
+
+from parse_bench.inference.providers.base import (
+    Provider,
+    ProviderConfigError,
+    ProviderPermanentError,
+    ProviderTransientError,
+)
+from parse_bench.inference.providers.parse._layout_utils import build_layout_pages
+from parse_bench.inference.providers.registry import register_provider
+from parse_bench.schemas.parse_output import ParseOutput
+from parse_bench.schemas.pipeline import PipelineSpec
+from parse_bench.schemas.pipeline_io import (
+    InferenceRequest,
+    InferenceResult,
+    RawInferenceResult,
+)
+from parse_bench.schemas.product import ProductType
+
+
+@register_provider("deepseekocr2")
+class DeepSeekOCR2Provider(Provider):
+    """
+    Provider for DeepSeek-OCR-2 Modal server.
+
+    Configuration options:
+        - server_url (str, required): Modal server predict endpoint URL
+        - timeout (int, default=600): Request timeout in seconds
+        - dpi (int, default=150): DPI for PDF to image conversion
+    """
+
+    def __init__(self, provider_name: str, base_config: dict[str, Any] | None = None):
+        super().__init__(provider_name, base_config)
+
+        server_url = self.base_config.get("server_url") or os.getenv("DEEPSEEKOCR2_SERVER_URL")
+        if not server_url:
+            raise ProviderConfigError("DeepSeekOCR2 provider requires 'server_url' in config.")
+        self._server_url: str = server_url
+
+        self._timeout = self.base_config.get("timeout", 600)
+        self._dpi = self.base_config.get("dpi", 150)
+
+    def _pdf_to_images(self, pdf_path: Path) -> list[bytes]:
+        """Render every PDF page to PNG bytes in source order."""
+        try:
+            from pdf2image import convert_from_path
+
+            images = convert_from_path(pdf_path, dpi=self._dpi)
+            if not images:
+                raise ProviderPermanentError(f"No pages found in PDF: {pdf_path}")
+            encoded: list[bytes] = []
+            for image in images:
+                buf = io.BytesIO()
+                image.save(buf, format="PNG")
+                encoded.append(buf.getvalue())
+            return encoded
+        except ImportError as e:
+            raise ProviderPermanentError("pdf2image is required.") from e
+        except Exception as e:
+            if "pdf2image" in str(e).lower():
+                raise
+            raise ProviderPermanentError(f"Error converting PDF to image: {e}") from e
+
+    def _read_image(self, file_path: Path) -> bytes:
+        try:
+            return file_path.read_bytes()
+        except Exception as e:
+            raise ProviderPermanentError(f"Error reading image file: {e}") from e
+
+    async def _call_api(self, session: aiohttp.ClientSession, image_b64: str) -> dict[str, Any]:
+        api_url = self._server_url.rstrip("/")
+
+        payload: dict[str, str] = {"image_base64": image_b64}
+
+        async with session.post(
+            api_url,
+            json=payload,
+            headers={"Content-Type": "application/json"},
+            timeout=aiohttp.ClientTimeout(total=self._timeout),
+        ) as resp:
+            if resp.status != 200:
+                error_text = await resp.text()
+                if resp.status in (408, 502, 503, 504):
+                    raise ProviderTransientError(f"HTTP {resp.status}: {error_text[:200]}")
+                raise ProviderPermanentError(f"HTTP {resp.status}: {error_text[:200]}")
+
+            result: dict[str, Any] = await resp.json()
+            if result.get("status") == "error":
+                raise ProviderPermanentError(result.get("error", "Unknown error from API"))
+
+            content: str = result.get("markdown", "")
+            if not content:
+                raise ProviderPermanentError("Empty markdown response from API")
+            return result
+
+    async def _run_inference_async(self, image_bytes: bytes) -> dict[str, Any]:
+        image_b64 = base64.b64encode(image_bytes).decode()
+
+        async with aiohttp.ClientSession() as session:
+            result = await self._call_api(session, image_b64)
+
+        return {
+            "markdown": result.get("markdown", ""),
+            "grounding_items": result.get("grounding_items", []),
+            "image_width": result.get("image_width"),
+            "image_height": result.get("image_height"),
+            "_config": {
+                "server_url": self._server_url,
+                "dpi": self._dpi,
+            },
+        }
+
+    async def _run_inference_pages_async(self, pages: list[bytes]) -> dict[str, Any]:
+        """Run each input page in order, retaining the legacy one-page shape."""
+        results = [await self._run_inference_async(page) for page in pages]
+        first = results[0]
+        if len(results) == 1:
+            return first
+        merged = dict(first)
+        merged["page_results"] = results
+        return merged
+
+    def run_inference(self, pipeline: PipelineSpec, request: InferenceRequest) -> RawInferenceResult:
+        if request.product_type != ProductType.PARSE:
+            raise ProviderPermanentError(
+                f"DeepSeekOCR2Provider only supports PARSE product type, got {request.product_type}"
+            )
+
+        started_at = datetime.now()
+
+        file_path = Path(request.source_file_path)
+        if not file_path.exists():
+            raise ProviderPermanentError(f"Source file not found: {file_path}")
+
+        suffix = file_path.suffix.lower()
+        if suffix == ".pdf":
+            page_images = self._pdf_to_images(file_path)
+        elif suffix in (".png", ".jpg", ".jpeg", ".webp", ".tiff", ".bmp"):
+            page_images = [self._read_image(file_path)]
+        else:
+            raise ProviderPermanentError(
+                f"Unsupported file type: {suffix}. Supported: .pdf, .png, .jpg, .jpeg, .webp, .tiff, .bmp"
+            )
+
+        try:
+            raw_output = asyncio.run(self._run_inference_pages_async(page_images))
+            completed_at = datetime.now()
+            latency_ms = int((completed_at - started_at).total_seconds() * 1000)
+
+            return RawInferenceResult(
+                request=request,
+                pipeline=pipeline,
+                pipeline_name=pipeline.pipeline_name,
+                product_type=request.product_type,
+                raw_output=raw_output,
+                started_at=started_at,
+                completed_at=completed_at,
+                latency_in_ms=latency_ms,
+            )
+
+        except (ProviderPermanentError, ProviderTransientError):
+            raise
+
+        except Exception as e:
+            completed_at = datetime.now()
+            latency_ms = int((completed_at - started_at).total_seconds() * 1000)
+
+            error_msg = str(e)
+            if isinstance(e, asyncio.TimeoutError):
+                error_msg = f"Request timed out after {self._timeout} seconds"
+
+            return RawInferenceResult(
+                request=request,
+                pipeline=pipeline,
+                pipeline_name=pipeline.pipeline_name,
+                product_type=request.product_type,
+                raw_output={
+                    "markdown": "",
+                    "_error": error_msg,
+                    "_error_type": type(e).__name__,
+                    "_config": {
+                        "server_url": self._server_url,
+                        "dpi": self._dpi,
+                    },
+                },
+                started_at=started_at,
+                completed_at=completed_at,
+                latency_in_ms=latency_ms,
+            )
+
+    @staticmethod
+    def _close_unclosed_table_tags(content: str) -> str:
+        """Auto-close unclosed HTML table tags from truncated model output."""
+        opens = content.count("<table>")
+        closes = content.count("</table>")
+        if opens > closes:
+            # Close any unclosed row/cell tags, then close the table
+            if content.rstrip().endswith(">"):
+                pass  # last tag is already closed
+            else:
+                # Truncated mid-cell — close the cell and row
+                content += "</td></tr>"
+            content += "</table>" * (opens - closes)
+        return content
+
+    @staticmethod
+    def _promote_first_row_to_thead(content: str) -> str:
+        """Wrap the first <tr> of each <table> in <thead> and convert <td> to <th>.
+
+        The grounding model outputs all cells as <td>, never using <th>/<thead>.
+        This heuristic promotes the first row to a header row, matching how
+        markdown2 handles pipe tables and improving header metric scores.
+        """
+
+        def _promote_table(match: re.Match[str]) -> str:
+            table_html: str = match.group(0)
+            # Find the first <tr>...</tr>
+            first_tr = re.search(r"<tr>(.*?)</tr>", table_html, re.DOTALL)
+            if not first_tr:
+                return table_html
+            first_tr_full: str = first_tr.group(0)
+            first_tr_inner: str = first_tr.group(1)
+            # Convert <td> to <th> in the first row
+            header_inner = first_tr_inner.replace("<td>", "<th>").replace("</td>", "</th>")
+            # Also handle <td with attributes
+            header_inner = re.sub(r"<td(\s)", r"<th\1", header_inner)
+            header_inner = re.sub(r"</td>", "</th>", header_inner)
+            thead = f"<thead><tr>{header_inner}</tr></thead>"
+            # Replace first <tr> with <thead> block
+            table_html = table_html.replace(first_tr_full, thead, 1)
+            return table_html
+
+        return re.sub(r"<table>.*?</table>", _promote_table, content, flags=re.DOTALL)
+
+    @staticmethod
+    def _sanitize_html_attributes(markdown: str) -> str:
+        """Quote unquoted HTML attributes for XML-based metric parsers."""
+
+        def _quote_attrs(match: re.Match) -> str:
+            tag_text = match.group(0)
+            tag_text = re.sub(
+                r'(\w+)=([^\s"\'<>=]+)',
+                r'\1="\2"',
+                tag_text,
+            )
+            return tag_text
+
+        return re.sub(r"<[^>]+>", _quote_attrs, markdown)
+
+    @staticmethod
+    def _convert_md_tables_to_html(content: str) -> str:
+        """Convert markdown pipe tables to HTML for GriTS/TEDS metrics."""
+        import markdown2
+
+        lines = content.split("\n")
+        result_parts: list[str] = []
+        table_lines: list[str] = []
+        in_table = False
+
+        for line in lines:
+            is_table_line = "|" in line and line.strip().startswith("|")
+            if is_table_line:
+                if not in_table:
+                    in_table = True
+                    table_lines = [line]
+                else:
+                    table_lines.append(line)
+            else:
+                if in_table:
+                    if len(table_lines) >= 2:
+                        table_md = "\n".join(table_lines)
+                        html = markdown2.markdown(table_md, extras=["tables"]).strip()
+                        if "<table>" in html.lower():
+                            result_parts.append(html)
+                        else:
+                            result_parts.extend(table_lines)
+                    else:
+                        result_parts.extend(table_lines)
+                    table_lines = []
+                    in_table = False
+                result_parts.append(line)
+
+        if in_table and len(table_lines) >= 2:
+            table_md = "\n".join(table_lines)
+            html = markdown2.markdown(table_md, extras=["tables"]).strip()
+            if "<table>" in html.lower():
+                result_parts.append(html)
+            else:
+                result_parts.extend(table_lines)
+        elif in_table:
+            result_parts.extend(table_lines)
+
+        return "\n".join(result_parts)
+
+    # Grounding label aliases -> the names the shared _layout_utils.LABEL_MAP
+    # understands (title/table/caption/footnote already match; these three differ).
+    _LABEL_ALIASES: dict[str, str] = {
+        "header": "page-header",
+        "footer": "page-footer",
+        "image": "picture",
+    }
+
+    def normalize(self, raw_result: RawInferenceResult) -> InferenceResult:
+        if raw_result.product_type != ProductType.PARSE:
+            raise ProviderPermanentError(
+                f"DeepSeekOCR2Provider only supports PARSE product type, got {raw_result.product_type}"
+            )
+
+        page_results = raw_result.raw_output.get("page_results")
+        if not isinstance(page_results, list) or not page_results:
+            page_results = [raw_result.raw_output]
+
+        page_markdowns: list[str] = []
+        layout_pages = []
+        for page_number, page_raw in enumerate(page_results, start=1):
+            markdown = page_raw.get("markdown", "")
+            if markdown:
+                # Auto-close unclosed HTML table tags (model truncates at max_tokens)
+                markdown = self._close_unclosed_table_tags(markdown)
+                markdown = self._convert_md_tables_to_html(markdown)
+                # Promote first row to <thead>/<th> (model outputs all <td>)
+                markdown = self._promote_first_row_to_thead(markdown)
+                markdown = self._sanitize_html_attributes(markdown)
+                page_markdowns.append(markdown)
+
+            # Build each page's layout from its grounding items. Grounding bboxes
+            # are [x1,y1,x2,y2] on a 0-999 grid; aliases map labels onto LABEL_MAP.
+            grounding_items = page_raw.get("grounding_items", [])
+            layout_items = [
+                {
+                    "label": self._LABEL_ALIASES.get(str(gi.get("label", "")).lower(), gi.get("label", "")),
+                    "bbox": gi.get("bbox", []),
+                }
+                for gi in grounding_items
+            ]
+            layout_pages.extend(
+                build_layout_pages(
+                    items=layout_items,
+                    image_width=page_raw.get("image_width", 0),
+                    image_height=page_raw.get("image_height", 0),
+                    markdown=markdown,
+                    page_number=page_number,
+                )
+            )
+        markdown = "\n\n".join(page_markdowns)
+
+        output = ParseOutput(
+            task_type="parse",
+            example_id=raw_result.request.example_id,
+            pipeline_name=raw_result.pipeline_name,
+            pages=[],
+            markdown=markdown,
+            layout_pages=layout_pages,
+        )
+
+        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/docling.py

@@ -0,0 +1,281 @@
+"""Provider for Docling parse via a custom HTTP endpoint."""
+
+import base64
+import os
+from datetime import datetime
+from pathlib import Path
+from typing import Any
+
+import requests
+from docling_core.types.doc.document import DoclingDocument
+
+from parse_bench.inference.providers.base import (
+    Provider,
+    ProviderConfigError,
+    ProviderPermanentError,
+    ProviderRateLimitError,
+    ProviderTransientError,
+)
+from parse_bench.inference.providers.parse._docling_common import _build_docling_layout_pages
+from parse_bench.inference.providers.registry import register_provider
+from parse_bench.schemas.parse_output import (
+    PageIR,
+    ParseLayoutPageIR,
+    ParseOutput,
+)
+from parse_bench.schemas.pipeline import PipelineSpec
+from parse_bench.schemas.pipeline_io import (
+    InferenceRequest,
+    InferenceResult,
+    RawInferenceResult,
+)
+from parse_bench.schemas.product import ProductType
+
+
+@register_provider("docling_parse")
+class DoclingParseProvider(Provider):
+    """
+    Provider for Docling PDF parsing via a custom HTTP endpoint.
+
+    This provider sends PDFs to a Docling endpoint and returns markdown
+    with tables formatted as HTML.
+    """
+
+    def __init__(
+        self,
+        provider_name: str,
+        base_config: dict[str, Any] | None = None,
+    ):
+        """
+        Initialize the Docling parse provider.
+
+        Args:
+            provider_name: Name of the provider
+            base_config: Optional configuration with:
+                - `api_key`: Optional bearer token for the endpoint
+                - `hf_token`: Deprecated alias for `api_key`
+                - `endpoint_url`: Endpoint URL (required)
+                - `timeout`: Request timeout in seconds (default: 120)
+        """
+        super().__init__(provider_name, base_config)
+
+        # Optional bearer token; keep `hf_token` / `HF_TOKEN` as a deprecated
+        # fallback for backwards compatibility with the previous HF deployment.
+        self._api_key = (
+            self.base_config.get("api_key")
+            or self.base_config.get("hf_token")
+            or os.getenv("DOCLING_PARSE_API_KEY")
+            or os.getenv("HF_TOKEN")
+            or ""
+        )
+
+        # Get endpoint URL (from config or env var)
+        self._endpoint_url = self.base_config.get("endpoint_url") or os.getenv("DOCLING_PARSE_ENDPOINT_URL")
+        if not self._endpoint_url:
+            raise ProviderConfigError(
+                "Docling endpoint URL is required. "
+                "Set DOCLING_PARSE_ENDPOINT_URL environment variable or "
+                "pass endpoint_url in pipeline config."
+            )
+
+        # Get timeout (default 120 seconds - PDF processing can be slow)
+        self._timeout = self.base_config.get("timeout", 120)
+
+    def _call_endpoint(self, pdf_bytes: bytes) -> dict[str, Any]:
+        """
+        Call the Docling endpoint with PDF bytes.
+
+        Args:
+            pdf_bytes: Raw PDF file bytes
+
+        Returns:
+            Raw JSON response from endpoint
+
+        Raises:
+            ProviderError: For any API errors
+        """
+        headers = {"Content-Type": "application/json"}
+        if self._api_key:
+            headers["Authorization"] = f"Bearer {self._api_key}"
+
+        # Encode PDF as base64
+        pdf_base64 = base64.b64encode(pdf_bytes).decode("utf-8")
+
+        payload = {
+            "inputs": {
+                "pdf_base64": pdf_base64,
+            }
+        }
+
+        try:
+            response = requests.post(
+                self._endpoint_url,
+                headers=headers,
+                json=payload,
+                timeout=self._timeout,
+            )
+            response.raise_for_status()
+            result_json = response.json()
+            if isinstance(result_json, list):
+                if not result_json:
+                    raise ProviderPermanentError("Endpoint returned an empty list response.")
+                first_result = result_json[0]
+                if not isinstance(first_result, dict):
+                    raise ProviderPermanentError("Endpoint returned a list response with a non-dict payload.")
+                result = first_result
+            elif isinstance(result_json, dict):
+                result = result_json
+            else:
+                raise ProviderPermanentError(
+                    f"Endpoint returned unsupported response type: {type(result_json).__name__}"
+                )
+            return result
+
+        except requests.exceptions.Timeout as e:
+            raise ProviderTransientError(f"Request timed out: {e}") from e
+        except requests.exceptions.ConnectionError as e:
+            raise ProviderTransientError(f"Connection error: {e}") from e
+        except requests.exceptions.HTTPError as e:
+            status_code = e.response.status_code if e.response else None
+            if status_code == 429:
+                raise ProviderRateLimitError(f"Rate limit exceeded: {e}") from e
+            elif status_code and 500 <= status_code < 600:
+                raise ProviderTransientError(f"Server error ({status_code}): {e}") from e
+            elif status_code and 400 <= status_code < 500:
+                raise ProviderPermanentError(f"Client error ({status_code}): {e}") from e
+            else:
+                raise ProviderPermanentError(f"HTTP error: {e}") from e
+        except (ProviderPermanentError, ProviderTransientError, ProviderRateLimitError):
+            raise
+        except Exception as e:
+            raise ProviderPermanentError(f"Unexpected error calling endpoint: {e}") from e
+
+    def run_inference(self, pipeline: PipelineSpec, request: InferenceRequest) -> RawInferenceResult:
+        """
+        Run inference and return raw results.
+
+        Args:
+            pipeline: Pipeline specification
+            request: Inference request
+
+        Returns:
+            Raw inference result
+
+        Raises:
+            ProviderError: For any provider-related failures
+        """
+        if request.product_type != ProductType.PARSE:
+            raise ProviderPermanentError(
+                f"DoclingParseProvider only supports PARSE product type, got {request.product_type}"
+            )
+
+        started_at = datetime.now()
+
+        # Check if file exists
+        source_path = Path(request.source_file_path)
+        if not source_path.exists():
+            raise ProviderPermanentError(f"Source file not found: {source_path}")
+
+        try:
+            # Read PDF bytes
+            pdf_bytes = source_path.read_bytes()
+
+            # Call endpoint
+            raw_output = self._call_endpoint(pdf_bytes)
+
+            completed_at = datetime.now()
+            latency_ms = int((completed_at - started_at).total_seconds() * 1000)
+
+            return RawInferenceResult(
+                request=request,
+                pipeline=pipeline,
+                pipeline_name=pipeline.pipeline_name,
+                product_type=request.product_type,
+                raw_output=raw_output,
+                started_at=started_at,
+                completed_at=completed_at,
+                latency_in_ms=latency_ms,
+            )
+
+        except (ProviderPermanentError, ProviderTransientError, ProviderRateLimitError):
+            raise
+        except Exception as e:
+            raise ProviderPermanentError(f"Unexpected error during inference: {e}") from e
+
+    def normalize(self, raw_result: RawInferenceResult) -> InferenceResult:
+        """
+        Normalize raw inference result to produce ParseOutput.
+
+        Args:
+            raw_result: Raw inference result from run_inference()
+
+        Returns:
+            Inference result with ParseOutput
+
+        Raises:
+            ProviderError: For any normalization failures
+        """
+        if raw_result.product_type != ProductType.PARSE:
+            raise ProviderPermanentError(
+                f"DoclingParseProvider only supports PARSE product type, got {raw_result.product_type}"
+            )
+
+        # Extract pages from response
+        # Response format:
+        # {
+        #   "pages": [{"page": 1, "markdown": "..."}, ...],
+        #   "markdown": "...",
+        #   "docling_document": {...},
+        # }
+        raw_pages = raw_result.raw_output.get("pages", [])
+        full_markdown = raw_result.raw_output.get("markdown", "")
+        raw_docling_document = raw_result.raw_output.get("docling_document")
+
+        # Convert to PageIR list (0-indexed)
+        pages: list[PageIR] = []
+        for page_data in raw_pages:
+            # Docling uses 1-indexed pages, we use 0-indexed
+            page_number = page_data.get("page", 1)
+            page_index = page_number - 1 if page_number > 0 else 0
+            markdown = page_data.get("markdown", "")
+
+            pages.append(PageIR(page_index=page_index, markdown=markdown))
+
+        # Sort by page index
+        pages.sort(key=lambda p: p.page_index)
+
+        # If we have pages but no full markdown, concatenate
+        if pages and not full_markdown:
+            full_markdown = "\n\n".join(p.markdown for p in pages)
+
+        layout_pages: list[ParseLayoutPageIR] = []
+        if raw_docling_document is not None:
+            try:
+                docling_document = DoclingDocument.model_validate(raw_docling_document)
+            except Exception as e:
+                raise ProviderPermanentError(f"Failed to validate docling_document payload: {e}") from e
+
+            layout_pages = _build_docling_layout_pages(
+                doc=docling_document,
+                raw_pages=[page for page in raw_pages if isinstance(page, dict)],
+            )
+
+        output = ParseOutput(
+            task_type="parse",
+            example_id=raw_result.request.example_id,
+            pipeline_name=raw_result.pipeline_name,
+            pages=pages,
+            layout_pages=layout_pages,
+            markdown=full_markdown,
+        )
+
+        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,
+        )