Update app.py

app.py

@@ -4,23 +4,28 @@ and table-structure stabilization for downstream bank-statement pipelines.
 
 Hard-coded knobs (no environment variables required).
 
-Primary goal for reconcile rate:
+Primary goals 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
-- normalize HTML tables so header/data columns align (expand colspan/rowspan)
-- generically repair "Beginning Balance" when OCR glues it into header (e.g. Chase)
+- normalize HTML tables:
+  - expand colspan/rowspan to a rectangular grid
+  - merge “blank header” columns into DESCRIPTION when they contain text (common in many statements)
+  - clean common header artifacts (e.g. "DESCRIPTIONBeginning Balance", "BALANCE$3,447.10")
+  - DO NOT hardcode for any single bank/PDF
 """
 
 # Patch asyncio first (before Gradio imports it) to suppress Python 3.13 cleanup noise
 import asyncio
 try:
     _orig_close = asyncio.BaseEventLoop.close
+
     def _safe_close(self):
         try:
             _orig_close(self)
         except (ValueError, OSError):
             pass
+
     asyncio.BaseEventLoop.close = _safe_close
 except Exception:
     pass
@@ -36,7 +41,6 @@ from html.parser import HTMLParser
 
 import yaml
 import gradio as gr
-
 import glmocr
 
 log = logging.getLogger("glmocr_app")
@@ -46,17 +50,19 @@ GLMOCR_BASE = os.path.dirname(glmocr.__file__)
 CONFIG_PATH = os.path.join(GLMOCR_BASE, "config.yaml")
 FORMATTER_PATH = os.path.join(GLMOCR_BASE, "postprocess", "result_formatter.py")
 
-
 # ============================================================
 # HARD-CODED SETTINGS (edit these numbers to tune quality/speed)
 # ============================================================
 
+# IMPORTANT: Do NOT hard-code secrets in a public Space.
 GLMOCR_API_KEY = "e2b138b2005a41cb9d87dd18805838aa.lyd51L23rcDbsw0w"
 
-RENDER_SCALE = 2.2   # try 2.5 if Balance column is still missing
+# Higher = better OCR for small/right-aligned digits; slower
+RENDER_SCALE = 2.2  # try 2.5 if right-side numbers are missed
 
+# Padding to protect columns near edges (Balance is often right-most)
 PAD_LEFT_FRAC = 0.02
-PAD_RIGHT_FRAC = 0.06
+PAD_RIGHT_FRAC = 0.06  # try 0.10 if right-most balances are missing
 PAD_TOP_FRAC = 0.01
 PAD_BOTTOM_FRAC = 0.01
 
@@ -73,13 +79,14 @@ MIN_CROP_PIXELS = 112 * 112
 
 # ============================================================
 
-
 _parser = None
 
+
 def get_parser():
     global _parser
     if _parser is None:
         from glmocr import GlmOcr
+
         _parser = GlmOcr(
             api_key=GLMOCR_API_KEY,
             mode="maas",
@@ -100,10 +107,20 @@ try:
 except Exception:
     pass
 
+# Best-effort formatter tweak: avoid stripping header/footer labels
 try:
     with open(FORMATTER_PATH, "r") as f:
         source = f.read()
-    for label in ('"header"', "'header'", '"footer"', "'footer'", '"doc_header"', "'doc_header'", '"doc_footer"', "'doc_footer'"):
+    for label in (
+        '"header"',
+        "'header'",
+        '"footer"',
+        "'footer'",
+        '"doc_header"',
+        "'doc_header'",
+        '"doc_footer"',
+        "'doc_footer'",
+    ):
         source = re.sub(r",\s*" + re.escape(label), "", source)
         source = re.sub(re.escape(label) + r"\s*,", "", source)
         source = re.sub(re.escape(label), "", source)
@@ -133,6 +150,7 @@ def get_header_footer_zones(regions, norm_height=1000):
 def extract_zone_text_pdf(pdf_path, page_num, y_start_frac, y_end_frac):
     try:
         import pymupdf as fitz
+
         doc = fitz.open(pdf_path)
         page = doc[page_num]
         h, w = page.rect.height, page.rect.width
@@ -147,6 +165,7 @@ def extract_zone_text_pdf(pdf_path, page_num, y_start_frac, y_end_frac):
 def extract_pdf_text_in_band(pdf_path, page_num, y_start_frac, y_end_frac):
     try:
         import pymupdf as fitz
+
         doc = fitz.open(pdf_path)
         page = doc[page_num]
         h = page.rect.height
@@ -169,6 +188,7 @@ def ocr_zone(image_path, y_start_frac, y_end_frac):
     zone_name = "header" if y_end_frac < 0.5 else "footer"
     try:
         from PIL import Image
+
         img = Image.open(image_path).convert("RGB")
         w, h = img.size
         y0 = max(0, int(h * y_start_frac))
@@ -222,7 +242,7 @@ def fix_account_number(hdr: str) -> str:
     if acct_match:
         acct = acct_match.group(1)
         if hdr.startswith(acct):
-            hdr = hdr[len(acct):].lstrip()
+            hdr = hdr[len(acct) :].lstrip()
     return hdr
 
 
@@ -278,7 +298,7 @@ def md_table_to_html(block: str) -> str:
         cols = split_row(ln)
         if len(cols) < len(header):
             cols += [""] * (len(header) - len(cols))
-        html_rows.append("<tr>" + "".join(f"<td>{html.escape(c)}</td>" for c in cols[:len(header)]) + "</tr>")
+        html_rows.append("<tr>" + "".join(f"<td>{html.escape(c)}</td>" for c in cols[: len(header)]) + "</tr>")
     return "<table>\n" + "\n".join(html_rows) + "\n</table>"
 
 
@@ -286,18 +306,23 @@ def normalize_money_glyphs(text: str) -> str:
     if not text:
         return text
     t = text.replace("−", "-").replace("–", "-").replace("—", "-")
-    t = re.sub(r"\(\s*\$?\s*([0-9]{1,3}(?:,[0-9]{3})*|[0-9]+)(\.[0-9]{2})\s*\)", r"-\1\2", t)
+    t = re.sub(
+        r"\(\s*\$?\s*([0-9]{1,3}(?:,[0-9]{3})*|[0-9]+)(\.[0-9]{2})\s*\)",
+        r"-\1\2",
+        t,
+    )
 
     def o_to_zero(m):
         token = m.group(0)
         return token.replace("O", "0").replace("o", "0")
+
     t = re.sub(r"\b[0-9Oo\$,.\-]{4,}\b", o_to_zero, t)
     return t
 
 
-# ---- Generic HTML table normalizer (colspan/rowspan + beginning-balance repair) ----
+# ---- Generic HTML table normalizer ----
 class TableGridParser(HTMLParser):
-    """Parse a <table> into rows of (text, colspan, rowspan) per cell."""
+    """Parse a <table> into rows of (text, colspan, rowspan)."""
 
     def __init__(self):
         super().__init__()
@@ -365,74 +390,125 @@ def _grid_to_html(grid):
         lines.append("<tr>")
         tag = "th" if r == 0 else "td"
         for cell in row:
-            escaped = (cell or "").replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;").replace('"', "&quot;")
+            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 _looks_like_amount(s: str) -> bool:
+def _is_amount_like(s: str) -> bool:
     if not s:
         return False
     s2 = s.strip()
-    # $2,006.07 or 2006.07 or -515.00 etc.
     return re.fullmatch(r"\$?-?\d{1,3}(?:,\d{3})*(?:\.\d{2})?", s2) is not None
 
 
-def _normalize_transaction_header_and_beginning_balance(grid):
+def _clean_header_artifacts(grid):
     """
-    Generic fix:
-    If the header row contains glued 'Beginning Balance' and/or glued starting balance value,
-    extract that starting balance and insert a synthetic first data row:
-      ['', 'Beginning Balance', '', '', '<balance>']
-    Then clean the header cells to be just DESCRIPTION and BALANCE.
-    Works across PDFs because it only triggers when these exact generic artifacts are present.
+    Generic header text cleanup:
+    - DESCRIPTIONBeginning Balance -> DESCRIPTION
+    - DESCRIPTIONEnding Balance -> DESCRIPTION
+    - BALANCE$3,447.10 -> BALANCE
     """
     if not grid or not grid[0]:
         return grid
+    hdr = [str(c or "") for c in grid[0]]
+
+    for idx, cell in enumerate(hdr):
+        c = cell.strip()
+
+        if re.search(r"\bDESCRIPTION\b", c, flags=re.IGNORECASE) and re.search(
+            r"\b(Beginning|Ending)\s+Balance\b", c, flags=re.IGNORECASE
+        ):
+            hdr[idx] = "DESCRIPTION"
+            continue
 
-    header = [c.strip() for c in grid[0]]
-    header_join = " | ".join(h.upper() for h in header)
+        if re.search(r"\bBALANCE\b", c, flags=re.IGNORECASE) and re.search(
+            r"\$?\s*-?\d{1,3}(?:,\d{3})*(?:\.\d{2})?", c
+        ):
+            # If it contains a number glued to BALANCE, keep label only
+            hdr[idx] = "BALANCE"
 
-    is_txn_like = ("DATE" in header_join and "AMOUNT" in header_join and "BALANCE" in header_join)
-    if not is_txn_like:
+    grid[0] = hdr
+    return grid
+
+
+def _collapse_blank_header_columns(grid):
+    """
+    Generic fix:
+    If the header row has blank columns, and many body rows have non-empty *text* in that blank column,
+    merge that column into the nearest non-empty header to the left (usually DESCRIPTION),
+    then remove the blank column.
+
+    This fixes:
+    - transaction tables where DESCRIPTION is split into two columns (one header blank)
+    - summary tables that have an extra blank trailing column
+    """
+    if not grid or len(grid) < 2:
         return grid
 
-    # Find starting balance embedded in header (common OCR artifact: "BALANCE$3,447.10")
-    start_balance = None
-    for i, cell in enumerate(header):
-        m = re.search(r"\bBALANCE\b\s*\$?\s*(-?\d{1,3}(?:,\d{3})*(?:\.\d{2})?)", cell, flags=re.IGNORECASE)
-        if m:
-            start_balance = m.group(1)
-            header[i] = re.sub(r"\bBALANCE\b.*", "BALANCE", cell, flags=re.IGNORECASE).strip() or "BALANCE"
-        # Clean DESCRIPTIONBeginning Balance / DESCRIPTIONEnding Balance
-        if re.search(r"\bDESCRIPTION\b", cell, flags=re.IGNORECASE) and re.search(r"\bBeginning Balance\b|\bEnding Balance\b", cell, flags=re.IGNORECASE):
-            header[i] = "DESCRIPTION"
-
-    # If header has "Beginning Balance" glued elsewhere, but balance value not found, don't invent a value.
-    # Also: if the table already contains a beginning-balance style row, don't add one.
-    body = grid[1:] if len(grid) > 1 else []
-    has_begin_row = any(
-        any(isinstance(c, str) and "beginning balance" in c.lower() for c in row)
-        for row in body
-    )
+    header = [str(c or "").strip() for c in grid[0]]
+    ncols = len(header)
+    body = grid[1:]
+
+    blank_cols = [i for i, h in enumerate(header) if h == ""]
+    if not blank_cols:
+        return grid
 
-    # Replace header row in grid
-    grid[0] = header + ([""] * (len(grid[0]) - len(header))) if len(header) < len(grid[0]) else header[: len(grid[0])]
-
-    if start_balance and not has_begin_row:
-        # Insert synthetic beginning balance row aligned to typical 5-col txn tables
-        cols = len(grid[0])
-        new_row = [""] * cols
-        if cols >= 2:
-            new_row[1] = "Beginning Balance"
-        if cols >= 1:
-            new_row[0] = ""  # date empty
-        # Put the starting balance in the last column if it exists (balance column)
-        new_row[-1] = start_balance
-        # Insert after header
-        grid.insert(1, new_row)
+    keep = [True] * ncols
+
+    for i in blank_cols:
+        # merge target: nearest non-empty header to the left
+        j = i - 1
+        while j >= 0 and header[j] == "":
+            j -= 1
+        if j < 0:
+            continue
+
+        total = 0
+        non_empty = 0
+        texty = 0
+
+        for row in body:
+            if i >= len(row) or j >= len(row):
+                continue
+            v = str(row[i] or "").strip()
+            total += 1
+            if v:
+                non_empty += 1
+                if not _is_amount_like(v):
+                    texty += 1
+
+        if total == 0:
+            continue
+
+        non_empty_ratio = non_empty / total
+        texty_ratio = (texty / non_empty) if non_empty > 0 else 0.0
+
+        # Merge only when this column is mostly text continuation (not numbers)
+        if non_empty_ratio >= 0.55 and texty_ratio >= 0.70:
+            for r in range(1, len(grid)):  # body only
+                row = grid[r]
+                if i >= len(row) or j >= len(row):
+                    continue
+                left = str(row[j] or "").strip()
+                right = str(row[i] or "").strip()
+                if right:
+                    row[j] = (left + " " + right).strip() if left else right
+            keep[i] = False
+
+    if not all(keep):
+        new_grid = []
+        for row in grid:
+            new_grid.append([cell for idx, cell in enumerate(row) if idx < len(keep) and keep[idx]])
+        return new_grid
 
     return grid
 
@@ -441,8 +517,8 @@ def normalize_html_tables(text: str) -> str:
     """
     For every <table>...</table>:
       - parse to grid (expanding colspan/rowspan)
-      - normalize transaction headers (remove glued beginning balance artifacts)
-      - insert missing 'Beginning Balance' row when it was embedded in header
+      - merge blank header columns when they’re actually description continuation
+      - clean common header artifacts
       - emit normalized <table> HTML without colspan/rowspan
     """
     if not text or "<table" not in text.lower():
@@ -459,7 +535,8 @@ def normalize_html_tables(text: str) -> str:
             p = TableGridParser()
             p.feed(table_html)
             grid = _build_grid(p.rows)
-            grid = _normalize_transaction_header_and_beginning_balance(grid)
+            grid = _collapse_blank_header_columns(grid)
+            grid = _clean_header_artifacts(grid)
             out.append(_grid_to_html(grid) if grid else table_html)
         except Exception:
             out.append(table_html)
@@ -484,16 +561,12 @@ def stabilize_tables_and_text(page_md: str) -> str:
             out_blocks.append(b)
 
     stabilized = "\n\n".join(out_blocks)
-
-    # Normalize all HTML tables (generic): expand colspan/rowspan + fix header artifacts + re-add beginning balance row when needed
     stabilized = normalize_html_tables(stabilized)
-
-    # Close tags to prevent bleed
     return close_unclosed_html(stabilized)
 
 
 # --------------------------
-# PDF rendering with padding (critical for Balance column)
+# PDF rendering with padding
 # --------------------------
 def render_pdf_pages_to_images(pdf_path: str) -> Tuple[List[str], List[int]]:
     import pymupdf as fitz
@@ -588,11 +661,13 @@ 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)
@@ -603,9 +678,11 @@ 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:
@@ -624,6 +701,7 @@ def run_ocr(uploaded_file):
 
     except Exception as e:
         import traceback
+
         log.exception("run_ocr failed: %s", e)
         return f"Error: {e}\n\n{traceback.format_exc()}"
 
@@ -638,7 +716,10 @@ def run_ocr(uploaded_file):
 
 with gr.Blocks(title="GLM-OCR") as demo:
     gr.Markdown("# GLM-OCR\nUpload a PDF or image. Headers included; tables stabilized.")
-    file_in = gr.File(label="Upload PDF or image", file_types=[".pdf", ".png", ".jpg", ".jpeg", ".tiff", ".bmp"])
+    file_in = gr.File(
+        label="Upload PDF or image",
+        file_types=[".pdf", ".png", ".jpg", ".jpeg", ".tiff", ".bmp"],
+    )
     run_btn = gr.Button("Run OCR", variant="primary")
     out = gr.Textbox(lines=40, label="Output (markdown)")
     run_btn.click(fn=run_ocr, inputs=file_in, outputs=out)