idp_types.py

# GENERATED by hub/build_hub_package.py from src/bodhan_genai/ocr/engine/types.py -- do not edit.
# Vendored so this repo is self-contained: `pip install transformers torch pillow` is the
# whole install. See indic_doc_parser.py for usage.

"""Plain-data types and configuration.

Every tunable that used to be an environment variable read at import time lives here instead.
That fixes a real defect: callers corrected the old module-level defaults by setting os.environ
*before* importing the engine, two shipped callers set different values, and importing it
directly gave a third pipeline. These defaults ARE the canonical recipe.

stdlib-only -- importable with no GPU stack and no PIL.
"""

from __future__ import annotations

import dataclasses
from dataclasses import dataclass, field

from idp_contract import TableFormat


class LayoutSchemaError(ValueError):
    """An incoming layout record cannot be interpreted unambiguously.

    Raised rather than warned because the alternative is silent and wrong. A label the taxonomy
    does not contain falls through to ``Text``, so a table would be sent the prose prompt and
    come back as flattened text with nothing in the output saying so. Our own detector only ever
    emits ``layout.labels.CLASSES``, so an unrecognised label always means either a third-party
    taxonomy -- which should declare ``type`` -- or a typo. Both are worth stopping for.

    ``map_label`` itself stays total: 19 of the 37 classes are deliberately absent from
    ``LABEL_TO_TYPE`` and correctly mean ``Text``, so leniency belongs there. Validation belongs
    here, at the boundary where foreign JSON enters.
    """


def _unknown(value, candidates, kind: str) -> str:
    """``unknown label 'Tabel'; did you mean 'Table'?`` -- the suggestion is most of the value."""
    import difflib

    match = difflib.get_close_matches(str(value), list(candidates), n=1, cutoff=0.6)
    hint = f"; did you mean {match[0]!r}?" if match else ""
    return f"unknown {kind} {value!r}{hint}"


@dataclass
class Block:
    """One detected region. ``text`` is None before OCR, and "" for blocks deliberately not
    transcribed (kept in place rather than deleted)."""

    order: int
    label: str
    type: str
    bbox_xyxy: list[float]
    conf: float
    text: str | None = None

    def as_record(self) -> dict:
        # Key order is load-bearing: json.dump writes insertion order and the regression gate
        # compares byte for byte. Do not reorder.
        record: dict = {
            "order": self.order,
            "label": self.label,
            "type": self.type,
            "bbox_xyxy": [round(float(v), 1) for v in self.bbox_xyxy],
            "conf": round(float(self.conf), 3),
        }
        if self.text is not None:
            record["text"] = self.text
        return record

    @classmethod
    def problems(cls, record: dict) -> list[str]:
        """Everything wrong with one incoming block record, as readable strings.

        Returns rather than raises so a whole page can be reported at once. Fixing a 60-block
        layout one exception at a time is the difference between one edit and sixty.
        """
        from idp_model_labels import CLASSES
        from idp_contract import DROP_TYPES, KEPT_BLOCK_TYPES

        found: list[str] = []
        for key in ("order", "bbox_xyxy"):
            if key not in record:
                found.append(f"missing required key {key!r}")

        bbox = record.get("bbox_xyxy")
        if bbox is not None:
            try:
                if len([float(v) for v in bbox]) != 4:
                    found.append(f"bbox_xyxy must be 4 numbers [x0, y0, x1, y1], got {bbox!r}")
            except (TypeError, ValueError):
                found.append(f"bbox_xyxy must be 4 numbers [x0, y0, x1, y1], got {bbox!r}")

        # An explicit `type` is the documented escape hatch for a foreign taxonomy, so the label
        # is not checked when one is given -- but the type itself is. Left unchecked, `Tabel`
        # is accepted verbatim and never matches `Table`, which is the same silent failure one
        # level up.
        declared = record.get("type")
        if declared:
            valid = tuple(KEPT_BLOCK_TYPES) + tuple(sorted(DROP_TYPES))
            if str(declared) not in valid:
                found.append(_unknown(declared, valid, "type"))
        else:
            label = str(record.get("label", "")).strip()
            if label.lower() not in {c.strip().lower() for c in CLASSES}:
                found.append(
                    _unknown(label, CLASSES, "label")
                    + ' -- use a spelling from layout.labels.CLASSES, or declare "type"'
                    " explicitly if your detector has its own taxonomy"
                )
        return found

    @classmethod
    def from_record(cls, record: dict, *, strict: bool = True) -> Block:
        # `type` is derived when absent, so a third-party layout carrying only labels works.
        from idp_contract import map_label

        if strict:
            found = cls.problems(record)
            if found:
                raise LayoutSchemaError("; ".join(found))

        label = record.get("label", "")
        return cls(
            order=int(record["order"]),
            label=str(label),
            type=str(record.get("type") or map_label(label)),
            bbox_xyxy=[float(v) for v in record["bbox_xyxy"]],
            conf=float(record.get("conf", 1.0)),
            text=record.get("text"),
        )

    def copy(self) -> Block:
        return dataclasses.replace(self, bbox_xyxy=list(self.bbox_xyxy))


@dataclass
class PageResult:
    image: str
    width: int
    height: int
    blocks: list[Block] = field(default_factory=list)
    markdown: str | None = None

    def as_record(self) -> dict:
        return {
            "image": self.image,
            "width": self.width,
            "height": self.height,
            "blocks": [b.as_record() for b in self.blocks],
        }

    @classmethod
    def from_record(cls, record: dict, *, strict: bool = True) -> PageResult:
        """``strict=False`` replays a layout without validating it -- for reading back files
        written before validation existed, not for new integrations."""
        blocks = list(record.get("blocks", []))

        if strict:
            errors = [
                f"  block[{i}] (order={b.get('order', '?')!r}): {problem}"
                for i, b in enumerate(blocks)
                for problem in Block.problems(b)
            ]
            if errors:
                raise LayoutSchemaError(
                    f"{len(errors)} problem(s) in layout for "
                    f"{record.get('image', '<unknown>')!r}:\n" + "\n".join(errors)
                )

        return cls(
            image=str(record["image"]),
            width=int(record["width"]),
            height=int(record["height"]),
            # Already validated above; re-checking each block would repeat the work.
            blocks=[Block.from_record(b, strict=False) for b in blocks],
        )


@dataclass(frozen=True)
class LayoutConfig:
    conf: float = 0.5  # below ~0.4, stains and page borders start scoring as blocks
    img_size: int = 1024
    device: str = "cuda"


@dataclass(frozen=True)
class CropConfig:
    """The clamp is on area, not on a side: pinning a side exploded elongated crops (a 122:1
    rule line became ~32k image tokens and wedged the engine)."""

    min_px_side: int = 256  # largest single measured win; 0 disables upscaling
    max_px_side: int = 1536  # token ceiling
    pad_px: int = 0

    @property
    def min_px(self) -> int:
        return self.min_px_side**2

    @property
    def max_px(self) -> int:
        return self.max_px_side**2


#: both -- text-like OR a larger Equation | text_only -- text-like only | eq_only -- Equation only
DEDUP_MODES = ("both", "text_only", "eq_only")


@dataclass(frozen=True)
class DedupConfig:
    """IndicDocLayout over-produces equation boxes nested inside the paragraphs and display
    arrays that already contain them; transcribing both emits the same math twice."""

    nest: bool = True
    mode: str = "both"
    contain: float = 0.90  # duplicate threshold in clean_layout
    wrap: float = 0.5  # a header counts as occupied at this containment
    nested: float = 0.70  # nested-equation threshold

    def __post_init__(self) -> None:
        if self.mode not in DEDUP_MODES:
            raise ValueError(f"DedupConfig.mode must be one of {DEDUP_MODES}, got {self.mode!r}")


@dataclass(frozen=True)
class RecognizerConfig:
    max_model_len: int = 8192
    max_tokens: int = 2048
    temperature: float = 0.0  # greedy: the only setting reproducible run to run
    gpu_memory_utilization: float = 0.80
    dtype: str = "bfloat16"
    # One giant batch over tens of thousands of multi-modal requests wedges the vLLM V1
    # scheduler at 100% util with no progress; ~2k chunks run clean.
    batch_size: int = 2048
    enforce_eager: bool = True  # skips ~4 min of torch.compile on a 0.8B model
    table_format: TableFormat = TableFormat.HTML

    def merged(self, **overrides) -> RecognizerConfig:
        """New config with non-None overrides applied; unknown keys raise."""
        known = {f.name for f in dataclasses.fields(self)}
        unknown = sorted(set(overrides) - known)
        if unknown:
            raise TypeError(f"Unknown RecognizerConfig field(s): {', '.join(unknown)}")
        return dataclasses.replace(self, **{k: v for k, v in overrides.items() if v is not None})