postprocess.py

"""Turn model text into the Markdown the OCR SDK would emit.

Inline ``\\(...\\)`` becomes ``$...$``. Display formulas lose ``\\[`` / ``\\]``
and are wrapped in ``$$`` later. Formula blocks then go through the same
LaTeX repair chain as the SDK. Table and chart OTSL becomes an HTML table.
"""

from __future__ import annotations

import html
import itertools
import re

from latex_repair import repair_latex_formula

INLINE_LATEX_RE = re.compile(r"\\\((.+?)\\\)")
DISPLAY_BRACKET_SEP_RE = re.compile(re.escape(r"\]") + r"\s*" + re.escape(r"\["))
OTSL_LABELS = {"table", "chart"}
DISPLAY_LABELS = {"display_formula", "formula"}

_NL = "<nl>"
_FCEL = "<fcel>"
_ECEL = "<ecel>"
_LCEL = "<lcel>"
_UCEL = "<ucel>"
_XCEL = "<xcel>"
_TAG_GROUP = "(?:<fcel>|<ecel>|<nl>|<lcel>|<ucel>|<xcel>)"
_FIND_CELLS = re.compile(f"{_TAG_GROUP}.*?(?={_TAG_GROUP}|$)", flags=re.DOTALL)
_ALL_TAGS = {_NL, _FCEL, _ECEL, _LCEL, _UCEL, _XCEL}
_TOKEN_RE = re.compile("(" + "|".join([_NL, _FCEL, _ECEL, _LCEL, _UCEL, _XCEL]) + ")")


class _Cell:
    def __init__(self, text, row_span, col_span, start_row, end_row, start_col, end_col):
        self.text = text
        self.row_span = row_span
        self.col_span = col_span
        self.start_row = start_row
        self.end_row = end_row
        self.start_col = start_col
        self.end_col = end_col


def _pad_rows(otsl: str) -> str:
    otsl = otsl.strip()
    if _NL not in otsl:
        return otsl + _NL
    rows = []
    for line in otsl.split(_NL):
        if not line:
            continue
        cells = _FIND_CELLS.findall(line)
        if not cells:
            continue
        filled = 0
        for index, cell in enumerate(cells):
            if cell.startswith(_FCEL):
                filled = index + 1
        rows.append((cells, len(cells), filled))
    if not rows:
        return _NL
    min_width = max(filled for _, _, filled in rows)
    max_width = max(total for _, total, _ in rows)
    best_cost = float("inf")
    width = max(min_width, max_width)
    for candidate in range(min_width, max(min_width, max_width) + 1):
        cost = sum(abs(total - candidate) for _, total, _ in rows)
        if cost < best_cost:
            best_cost = cost
            width = candidate
    repaired = []
    for cells, total, _filled in rows:
        if total > width:
            cells = cells[:width]
        else:
            cells = cells + [_ECEL] * (width - total)
        repaired.append("".join(cells))
    return _NL.join(repaired) + _NL


def _tokens_and_text(source: str):
    tokens = _TOKEN_RE.findall(source)
    parts = [part for part in _TOKEN_RE.split(source) if part.strip()]
    return tokens, parts


def _cells(texts, tokens):
    rows = [list(group) for is_break, group in itertools.groupby(tokens, lambda token: token == _NL) if not is_break]
    if rows:
        width = max(len(row) for row in rows)
        for row in rows:
            while len(row) < width:
                row.append(_ECEL)
        aligned = []
        cursor = 0
        for row in rows:
            for token in row:
                aligned.append(token)
                if cursor < len(texts) and texts[cursor] == token:
                    cursor += 1
                    if cursor < len(texts) and texts[cursor] not in _ALL_TAGS:
                        aligned.append(texts[cursor])
                        cursor += 1
            aligned.append(_NL)
            if cursor < len(texts) and texts[cursor] == _NL:
                cursor += 1
        texts = aligned

    def count_right(column, row, kinds):
        count = 0
        while column < len(rows[row]) and rows[row][column] in kinds:
            column += 1
            count += 1
        return count

    def count_down(column, row, kinds):
        count = 0
        while row < len(rows) and column < len(rows[row]) and rows[row][column] in kinds:
            row += 1
            count += 1
        return count

    cells = []
    row_index = 0
    column_index = 0
    for index, text in enumerate(texts):
        if text in {_FCEL, _ECEL}:
            content = ""
            offset = 1
            if text != _ECEL:
                content = texts[index + 1] if index + 1 < len(texts) else ""
                offset = 2
            right = texts[index + offset] if index + offset < len(texts) else ""
            below = ""
            if row_index + 1 < len(rows) and column_index < len(rows[row_index + 1]):
                below = rows[row_index + 1][column_index]
            col_span = 1
            row_span = 1
            if right in {_LCEL, _XCEL}:
                col_span += count_right(column_index + 1, row_index, {_LCEL, _XCEL})
            if below in {_UCEL, _XCEL}:
                row_span += count_down(column_index, row_index + 1, {_UCEL, _XCEL})
            cells.append(
                _Cell(
                    content.strip(),
                    row_span,
                    col_span,
                    row_index,
                    row_index + row_span,
                    column_index,
                    column_index + col_span,
                )
            )
        if text in _ALL_TAGS and text != _NL:
            column_index += 1
        if text == _NL:
            row_index += 1
            column_index = 0
    return cells, rows


def convert_otsl_to_html(otsl_content: str) -> str:
    padded = _pad_rows(html.unescape(otsl_content))
    tokens, texts = _tokens_and_text(padded)
    cells, rows = _cells(texts, tokens)
    if not cells:
        return ""
    row_count = len(rows)
    col_count = max((len(row) for row in rows), default=0)
    grid = [[None for _ in range(col_count)] for _ in range(row_count)]
    for cell in cells:
        for row_index in range(min(cell.start_row, row_count), min(cell.end_row, row_count)):
            for col_index in range(min(cell.start_col, col_count), min(cell.end_col, col_count)):
                grid[row_index][col_index] = cell
    body = []
    for row_index in range(row_count):
        parts = ["<tr>"]
        for col_index in range(col_count):
            cell = grid[row_index][col_index]
            if cell is None:
                parts.append("<td></td>")
                continue
            if cell.start_row != row_index or cell.start_col != col_index:
                continue
            attrs = ""
            if cell.row_span > 1:
                attrs += f' rowspan="{cell.row_span}"'
            if cell.col_span > 1:
                attrs += f' colspan="{cell.col_span}"'
            parts.append(f"<td{attrs}>{html.escape(cell.text.strip())}</td>")
        parts.append("</tr>")
        body.append("".join(parts))
    return "<table>" + "".join(body) + "</table>"


def _strip_wrapping_delimiter(content: str) -> str:
    for left, right in (("$$", "$$"), ("$", "$"), ("\\(", "\\)"), ("\\[", "\\]")):
        if content.startswith(left) and content.endswith(right) and len(content) > len(left) + len(right):
            return content[len(left) : len(content) - len(right)].strip()
    return content


def _strip_html_fence(content: str) -> str:
    if content.startswith("```html"):
        content = content[7:]
    elif content.startswith("```"):
        content = content[3:]
    if content.endswith("```"):
        content = content[:-3]
    content = content.strip()
    opened = content.count("<table")
    closed = content.count("</table>")
    if opened > closed:
        content += "</table>" * (opened - closed)
    return content


def normalize_content(text: str, label: str) -> str:
    out = (text or "").strip()
    if label in OTSL_LABELS:
        out = _strip_html_fence(out)
    elif label in DISPLAY_LABELS or label == "inline_formula":
        out = _strip_wrapping_delimiter(out)
    if label in DISPLAY_LABELS:
        out = DISPLAY_BRACKET_SEP_RE.sub("$$   $$", out)
        out = out.replace("\\[", "").replace("\\]", "")
        out = repair_latex_formula(out, label="display_formula")
    elif label == "inline_formula":
        out = repair_latex_formula(out, label="inline_formula")
        out = INLINE_LATEX_RE.sub(lambda match: f"${match.group(1).strip()}$", out)
    else:
        out = INLINE_LATEX_RE.sub(lambda match: f"${match.group(1).strip()}$", out)
    out = re.sub(r"[ \t]+", " ", out)
    out = re.sub(r"\n{3,}", "\n\n", out)
    if label in OTSL_LABELS:
        table = convert_otsl_to_html(out)
        if table:
            return table
    return out


def format_block(label: str, text: str) -> str:
    content = normalize_content(text, label)
    if label == "doc_title":
        return f"# {content}"
    if label == "paragraph_title":
        return f"## {content}"
    if label == "figure_title":
        return f"*{content}*"
    if label in DISPLAY_LABELS:
        return f"$$\n{content}\n$$"
    if label == "inline_formula":
        return f"${content}$"
    if label == "image":
        return "![image]()"
    return content