Spaces:
Running on Zero
Running on Zero
| """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 | |