Copy from bodhan-ai/indic-doc-parser

7b2177e verified about 2 months ago

4.47 kB

# GENERATED by hub/build_hub_package.py from src/bodhan_genai/ocr/templates/contract.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.

"""IndicDocParser's contract: prompts, block taxonomy, output schema.
Two vocabularies meet here and mixing them fails silently. **Labels** are what IndicDocLayout
emits (``ocr.layout.labels.CLASSES``); MARGINALIA and HEAD_FOOT match them case-SENSITIVELY.
**Types** are the coarse categories a label maps to, and select the prompt. OCR_SKIP_LABELS is
the exception: matched case-insensitively.
stdlib-only -- importable with no GPU stack and no PIL.
"""

from __future__ import annotations

from enum import StrEnum


class TableFormat(StrEnum):
    """HTML is the default: colspan/rowspan and in-cell breaks have no GFM spelling, so a
    merged-cell table rendered as markdown silently loses its structure."""

    HTML = "html"
    MARKDOWN = "markdown"


TEXT_PROMPT = (
    "Transcribe the text in this image. Write any mathematical expressions in LaTeX, "
    "using $...$ for inline math and $$...$$ for display equations."
)
EQUATION_PROMPT = "Output only the LaTeX for this equation image."
TABLE_PROMPTS = {
    TableFormat.HTML: (
        "Convert this table image to HTML. Preserve the structure exactly, using colspan and "
        "rowspan for merged cells and <br/> for line breaks within a cell. "
        "Output only the HTML table."
    ),
    TableFormat.MARKDOWN: (
        "Convert this table image to a GitHub-flavored markdown table. Output only the table."
    ),
}


def prompt_for(block_type: str, table_format: TableFormat = TableFormat.HTML) -> str:
    if block_type == "Table":
        return TABLE_PROMPTS[TableFormat(table_format)]
    if block_type == "Equation":
        return EQUATION_PROMPT
    return TEXT_PROMPT


# Labels absent from this map fall through to "Text" by design -- Question, Paragraph, Answer,
# List, MCQ, Code, Reference and the rest all carry prose.
LABEL_TO_TYPE = {
    "table": "Table",
    "table-caption": "Caption",
    "equation": "Equation",
    "expression": "Equation",
    "diagram": "Figure",
    "chart": "Figure",
    "image": "Picture",
    "image-caption": "Caption",
    "title": "Title",
    "chapter-title": "Title",
    "section-title": "SectionHeader",
    "sub-section-title": "SectionHeader",
    "sub-sub-section-title": "SectionHeader",
    "header": "PageHeader",
    "footer": "PageFooter",
    "page-number": "PageNumber",
    "folio": "PageNumber",
    "footnote": "Footnote",
}


def map_label(label) -> str:
    """Label -> pipeline type. Unknown labels are ``Text``."""
    return LABEL_TO_TYPE.get(str(label).strip().lower(), "Text")


# Exactly the types map_label can produce, less DROP_TYPES. test_contract asserts this, so a
# dead or undocumented type cannot creep in.
KEPT_BLOCK_TYPES = (
    "Text",
    "Title",
    "SectionHeader",
    "Table",
    "Equation",
    "Caption",
    "Footnote",
    "PageHeader",
    "PageFooter",
    "PageNumber",
)

# Never cropped, never reconstructed: pictorial regions have no text and invite hallucination.
# Deliberately narrow -- page numbers and margin text DO reach the recognizer.
DROP_TYPES = frozenset({"Figure", "Picture"})

# Cleaned as their own group so a page-spanning paragraph cannot swallow a page number.
MARGINALIA = frozenset({"Header", "Footer", "Page-number", "Folio"})
HEAD_FOOT = ("Header", "Footer")

# Never sent to the recognizer, but NOT deleted: these keep their place in the output with
# text "", so a consumer can still see what was detected and where.
OCR_SKIP_LABELS = frozenset({"header", "footer", "diagram", "image", "chart", "advertisement"})


def is_transcribed(label) -> bool:
    return str(label).strip().lower() not in OCR_SKIP_LABELS


OUTPUT_BLOCK_SCHEMA = {
    "order": "int    -- reading-order rank (0 = first)",
    "label": "str    -- raw IndicDocLayout label",
    "type": "str    -- pipeline type (see KEPT_BLOCK_TYPES)",
    "bbox_xyxy": "[float x4] -- pixel [x0, y0, x1, y1], clamped to the page",
    "conf": "float  -- detection confidence",
    "text": "str    -- transcription; '' when not sent to the recognizer",
}

OUTPUT_PAGE_SCHEMA = {
    "image": "str   -- source image filename",
    "width": "int   -- page width in pixels",
    "height": "int   -- page height in pixels",
    "blocks": "[BLOCK] -- in reading order",
}

	# GENERATED by hub/build_hub_package.py from src/bodhan_genai/ocr/templates/contract.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.

	"""IndicDocParser's contract: prompts, block taxonomy, output schema.

	Two vocabularies meet here and mixing them fails silently. Labels are what IndicDocLayout
	emits (``ocr.layout.labels.CLASSES``); MARGINALIA and HEAD_FOOT match them case-SENSITIVELY.
	Types are the coarse categories a label maps to, and select the prompt. OCR_SKIP_LABELS is
	the exception: matched case-insensitively.

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

	from __future__ import annotations

	from enum import StrEnum


	class TableFormat(StrEnum):
	"""HTML is the default: colspan/rowspan and in-cell breaks have no GFM spelling, so a
	merged-cell table rendered as markdown silently loses its structure."""

	HTML = "html"
	MARKDOWN = "markdown"


	TEXT_PROMPT = (
	"Transcribe the text in this image. Write any mathematical expressions in LaTeX, "
	"using $...$ for inline math and $$...$$ for display equations."
	)
	EQUATION_PROMPT = "Output only the LaTeX for this equation image."
	TABLE_PROMPTS = {
	TableFormat.HTML: (
	"Convert this table image to HTML. Preserve the structure exactly, using colspan and "
	"rowspan for merged cells and <br/> for line breaks within a cell. "
	"Output only the HTML table."
	),
	TableFormat.MARKDOWN: (
	"Convert this table image to a GitHub-flavored markdown table. Output only the table."
	),
	}


	def prompt_for(block_type: str, table_format: TableFormat = TableFormat.HTML) -> str:
	if block_type == "Table":
	return TABLE_PROMPTS[TableFormat(table_format)]
	if block_type == "Equation":
	return EQUATION_PROMPT
	return TEXT_PROMPT


	# Labels absent from this map fall through to "Text" by design -- Question, Paragraph, Answer,
	# List, MCQ, Code, Reference and the rest all carry prose.
	LABEL_TO_TYPE = {
	"table": "Table",
	"table-caption": "Caption",
	"equation": "Equation",
	"expression": "Equation",
	"diagram": "Figure",
	"chart": "Figure",
	"image": "Picture",
	"image-caption": "Caption",
	"title": "Title",
	"chapter-title": "Title",
	"section-title": "SectionHeader",
	"sub-section-title": "SectionHeader",
	"sub-sub-section-title": "SectionHeader",
	"header": "PageHeader",
	"footer": "PageFooter",
	"page-number": "PageNumber",
	"folio": "PageNumber",
	"footnote": "Footnote",
	}


	def map_label(label) -> str:
	"""Label -> pipeline type. Unknown labels are ``Text``."""
	return LABEL_TO_TYPE.get(str(label).strip().lower(), "Text")


	# Exactly the types map_label can produce, less DROP_TYPES. test_contract asserts this, so a
	# dead or undocumented type cannot creep in.
	KEPT_BLOCK_TYPES = (
	"Text",
	"Title",
	"SectionHeader",
	"Table",
	"Equation",
	"Caption",
	"Footnote",
	"PageHeader",
	"PageFooter",
	"PageNumber",
	)

	# Never cropped, never reconstructed: pictorial regions have no text and invite hallucination.
	# Deliberately narrow -- page numbers and margin text DO reach the recognizer.
	DROP_TYPES = frozenset({"Figure", "Picture"})

	# Cleaned as their own group so a page-spanning paragraph cannot swallow a page number.
	MARGINALIA = frozenset({"Header", "Footer", "Page-number", "Folio"})
	HEAD_FOOT = ("Header", "Footer")

	# Never sent to the recognizer, but NOT deleted: these keep their place in the output with
	# text "", so a consumer can still see what was detected and where.
	OCR_SKIP_LABELS = frozenset({"header", "footer", "diagram", "image", "chart", "advertisement"})


	def is_transcribed(label) -> bool:
	return str(label).strip().lower() not in OCR_SKIP_LABELS


	OUTPUT_BLOCK_SCHEMA = {
	"order": "int -- reading-order rank (0 = first)",
	"label": "str -- raw IndicDocLayout label",
	"type": "str -- pipeline type (see KEPT_BLOCK_TYPES)",
	"bbox_xyxy": "[float x4] -- pixel [x0, y0, x1, y1], clamped to the page",
	"conf": "float -- detection confidence",
	"text": "str -- transcription; '' when not sent to the recognizer",
	}

	OUTPUT_PAGE_SCHEMA = {
	"image": "str -- source image filename",
	"width": "int -- page width in pixels",
	"height": "int -- page height in pixels",
	"blocks": "[BLOCK] -- in reading order",
	}