Update app.py

app.py

@@ -8,6 +8,7 @@ Primary goal for reconcile rate:
 - preserve right-most columns (often "Balance") by higher DPI render + right padding
 - keep tables as tables (convert markdown pipe tables -> HTML table)
 - return ---page-separator--- between pages
+- expand colspan/rowspan in HTML tables so header and data rows align (fixes reconciliation)
 """
 
 # Patch asyncio first (before Gradio imports it) to suppress Python 3.13 cleanup noise
@@ -28,6 +29,8 @@ import os
 import re
 import html
 import tempfile
+from collections import defaultdict
+from html.parser import HTMLParser
 from typing import List, Tuple
 
 import yaml
@@ -259,6 +262,116 @@ def close_unclosed_html(md: str) -> str:
     return md
 
 
+# ---- Colspan/rowspan expansion: align header and data rows for any PDF ----
+class TableGridParser(HTMLParser):
+    """Parse <table> HTML into a list of rows; each row is a list of (text, colspan, rowspan)."""
+
+    def __init__(self):
+        super().__init__()
+        self.rows = []
+        self._current_row = []
+        self._cell_text = []
+        self._colspan = 1
+        self._rowspan = 1
+
+    def handle_starttag(self, tag, attrs):
+        if tag == "tr":
+            self._current_row = []
+        elif tag in ("td", "th"):
+            attrs_d = dict(attrs)
+            self._colspan = max(1, int(attrs_d.get("colspan", 1)))
+            self._rowspan = max(1, int(attrs_d.get("rowspan", 1)))
+            self._cell_text = []
+
+    def handle_endtag(self, tag):
+        if tag in ("td", "th"):
+            text = "".join(self._cell_text).strip().replace("\n", " ")
+            self._current_row.append((text, self._colspan, self._rowspan))
+        elif tag == "tr":
+            self.rows.append(self._current_row)
+
+    def handle_data(self, data):
+        self._cell_text.append(data)
+
+
+def _build_grid(rows_data):
+    """
+    Build a 2D grid from rows_data (list of list of (text, colspan, rowspan)).
+    Expands colspan and rowspan so every row has the same number of columns.
+    """
+    if not rows_data:
+        return []
+    blocked = defaultdict(set)
+    grid = []
+
+    for r, row_cells in enumerate(rows_data):
+        grid.append([])
+        col = 0
+        for content, C, R in row_cells:
+            while col in blocked[r]:
+                grid[r].append("")
+                col += 1
+            for k in range(C):
+                grid[r].append(content if k == 0 else "")
+            for k in range(1, R):
+                blocked[r + k].add(col)
+            col += C
+
+    max_cols = max(len(row) for row in grid) if grid else 0
+    for row in grid:
+        while len(row) < max_cols:
+            row.append("")
+    return grid
+
+
+def _grid_to_html(grid):
+    """Turn a 2D grid of cell texts into a single <table> HTML string (no colspan/rowspan)."""
+    if not grid:
+        return ""
+    lines = ["<table>"]
+    for i, row in enumerate(grid):
+        lines.append("<tr>")
+        tag = "th" if i == 0 else "td"
+        for cell in row:
+            escaped = (cell or "").replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;").replace('"', "&quot;")
+            lines.append(f"<{tag}>{escaped}</{tag}>")
+        lines.append("</tr>")
+    lines.append("</table>")
+    return "\n".join(lines)
+
+
+def expand_colspan_rowspan_in_tables(text):
+    """
+    Find every <table>...</table> in text and replace with a version where
+    colspan and rowspan are expanded so all rows have the same number of cells.
+    Works for any PDF; prevents header/data misalignment that breaks reconciliation.
+    """
+    if not text or "<table" not in text.lower():
+        return text
+
+    pattern = re.compile(r"<table[^>]*>.*?</table>", re.DOTALL | re.IGNORECASE)
+    result = []
+    last_end = 0
+
+    for match in pattern.finditer(text):
+        result.append(text[last_end : match.start()])
+        table_html = match.group(0)
+        try:
+            parser = TableGridParser()
+            parser.feed(table_html)
+            if parser.rows:
+                grid = _build_grid(parser.rows)
+                result.append(_grid_to_html(grid))
+            else:
+                result.append(table_html)
+        except Exception:
+            result.append(table_html)
+        last_end = match.end()
+
+    result.append(text[last_end:])
+    return "".join(result)
+
+
 def looks_like_markdown_table(block: str) -> bool:
     """Detect simple markdown pipe tables."""
     lines = [ln.rstrip() for ln in block.strip().splitlines() if ln.strip()]
@@ -312,7 +425,7 @@ def normalize_money_glyphs(text: str) -> str:
 
 
 def stabilize_tables_and_text(page_md: str) -> str:
-    """Convert markdown pipe tables to HTML and close tags."""
+    """Convert markdown pipe tables to HTML, expand colspan/rowspan, and close tags."""
     if not page_md:
         return page_md
     page_md = normalize_money_glyphs(page_md)
@@ -326,6 +439,7 @@ def stabilize_tables_and_text(page_md: str) -> str:
             out_blocks.append(b)
 
     stabilized = "\n\n".join(out_blocks)
+    stabilized = expand_colspan_rowspan_in_tables(stabilized)
     return close_unclosed_html(stabilized)
 
 
@@ -425,13 +539,11 @@ def run_ocr(uploaded_file):
             he = header_end_frac if header_end_frac is not None else DEFAULT_ZONE_FRAC
             fs = footer_start_frac if footer_start_frac is not None else (1.0 - DEFAULT_ZONE_FRAC)
 
-            # clamp
             he = max(0.02, min(0.25, he))
             fs = max(0.75, min(0.98, fs))
 
             parts = []
 
-            # Header inclusion: PDF text -> band words -> OCR band
             hdr = ""
             if is_pdf:
                 hdr = extract_zone_text_pdf(path, page_num, 0, he)
@@ -442,11 +554,9 @@ def run_ocr(uploaded_file):
             if hdr and hdr.strip():
                 parts.append(fix_account_number(normalize_money_glyphs(hdr.strip())))
 
-            # Main OCR markdown, stabilized
             if page_md and page_md.strip():
                 parts.append(stabilize_tables_and_text(page_md.strip()))
 
-            # Optional footer
             if ENABLE_FOOTER_OCR and page_num < len(page_images):
                 ftr = ""
                 if is_pdf:
@@ -469,7 +579,6 @@ def run_ocr(uploaded_file):
         return f"Error: {e}\n\n{traceback.format_exc()}"
 
     finally:
-        # cleanup rendered images
         for p in page_images:
             try:
                 if isinstance(p, str) and p.endswith(".png") and "glmocr_page_" in os.path.basename(p):