Update app.py

app.py

@@ -9,7 +9,7 @@ Primary goal for reconcile rate:
 - 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)
-- clean common header artifacts (e.g. "DESCRIPTIONBeginning Balance", "BALANCE$3,447.10")
+- generically repair "Beginning Balance" when OCR glues it into header (e.g. Chase)
 """
 
 # Patch asyncio first (before Gradio imports it) to suppress Python 3.13 cleanup noise
@@ -51,39 +51,29 @@ FORMATTER_PATH = os.path.join(GLMOCR_BASE, "postprocess", "result_formatter.py")
 # HARD-CODED SETTINGS (edit these numbers to tune quality/speed)
 # ============================================================
 
-# 1) GLM-OCR MaaS API key
-# IMPORTANT: Do NOT hard-code secrets in a public Space.
-# If your Space is public, switch to HF Secrets instead.
 GLMOCR_API_KEY = "e2b138b2005a41cb9d87dd18805838aa.lyd51L23rcDbsw0w"
 
-# 2) Render quality (higher = better OCR for small/right-aligned digits; slower)
 RENDER_SCALE = 2.2   # try 2.5 if Balance column is still missing
 
-# 3) Add padding to protect columns near edges (Balance is usually right-most)
 PAD_LEFT_FRAC = 0.02
-PAD_RIGHT_FRAC = 0.06  # try 0.10 if right-most balances are missing
+PAD_RIGHT_FRAC = 0.06
 PAD_TOP_FRAC = 0.01
 PAD_BOTTOM_FRAC = 0.01
 
-# 4) Mild contrast boost (helps faint gray text)
 ENABLE_CONTRAST = True
 
-# 5) Header/footer band heuristics
-DEFAULT_ZONE_FRAC = 0.12      # OCR band for header (top 12%) when regions not available
-PDF_HEADER_BAND_FRAC = 0.10   # PDF text fallback: take top 10% words if clip returns empty
+DEFAULT_ZONE_FRAC = 0.12
+PDF_HEADER_BAND_FRAC = 0.10
 
-# Footer: disabled by default to avoid duplicating what GLM already returns
 ENABLE_FOOTER_OCR = False
-PDF_FOOTER_BAND_FRAC = 0.88   # bottom 12% (if footer enabled)
+PDF_FOOTER_BAND_FRAC = 0.88
 
-# MaaS minimum image sizes for crops; we pad if needed
 MIN_CROP_HEIGHT = 112
 MIN_CROP_PIXELS = 112 * 112
 
 # ============================================================
 
 
-# Single shared parser to avoid re-init per request
 _parser = None
 
 def get_parser():
@@ -110,7 +100,6 @@ try:
 except Exception:
     pass
 
-# Best-effort formatter tweak: avoid stripping header/footer labels
 try:
     with open(FORMATTER_PATH, "r") as f:
         source = f.read()
@@ -128,7 +117,6 @@ except Exception:
 # Header/footer helpers
 # --------------------------
 def get_header_footer_zones(regions, norm_height=1000):
-    """Infer header/footer extents from bbox regions if present."""
     if not regions:
         return None, None
     y_tops, y_bottoms = [], []
@@ -143,7 +131,6 @@ def get_header_footer_zones(regions, norm_height=1000):
 
 
 def extract_zone_text_pdf(pdf_path, page_num, y_start_frac, y_end_frac):
-    """Extract text from a horizontal band using a clip rect (works if PDF has text layer)."""
     try:
         import pymupdf as fitz
         doc = fitz.open(pdf_path)
@@ -158,7 +145,6 @@ 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):
-    """Extract words whose bbox intersects a vertical band (robust fallback)."""
     try:
         import pymupdf as fitz
         doc = fitz.open(pdf_path)
@@ -180,7 +166,6 @@ def extract_pdf_text_in_band(pdf_path, page_num, y_start_frac, y_end_frac):
 
 
 def ocr_zone(image_path, y_start_frac, y_end_frac):
-    """Run OCR on a horizontal band. Pads small crops to meet MaaS minimum size."""
     zone_name = "header" if y_end_frac < 0.5 else "footer"
     try:
         from PIL import Image
@@ -227,7 +212,6 @@ def ocr_zone(image_path, y_start_frac, y_end_frac):
 
 
 def fix_account_number(hdr: str) -> str:
-    """Fix common account-number formatting issues."""
     if not hdr:
         return hdr
     if "Account Number:" in hdr and "Account Number: " not in hdr:
@@ -246,7 +230,6 @@ def fix_account_number(hdr: str) -> str:
 # Table stabilization helpers
 # --------------------------
 def close_unclosed_html(md: str) -> str:
-    """Close unclosed <table>/<tr>/<td> tags to prevent bleed."""
     if not md:
         return md
     open_tags = re.findall(r"<(table|tbody|thead|tr|td|th)\b", md, flags=re.IGNORECASE)
@@ -263,46 +246,58 @@ def close_unclosed_html(md: str) -> str:
     return md
 
 
-def clean_table_header_artifacts(text: str) -> str:
-    """
-    Generic cleanup for common OCR artifacts in table headers:
-    - 'DESCRIPTIONBeginning Balance' -> 'DESCRIPTION'
-    - 'DESCRIPTIONEnding Balance' -> 'DESCRIPTION'
-    - 'BALANCE$3,447.10' -> 'BALANCE'
-    - 'BALANCE 3,447.10' -> 'BALANCE'
-    Works for any PDF; no bank-specific logic.
-    """
-    if not text:
-        return text
+def looks_like_markdown_table(block: str) -> bool:
+    lines = [ln.rstrip() for ln in block.strip().splitlines() if ln.strip()]
+    if len(lines) < 2:
+        return False
+    if "|" not in lines[0]:
+        return False
+    sep = lines[1].replace(" ", "")
+    return ("---" in sep) and ("|" in sep)
 
-    # DESCRIPTION + (Beginning/Ending Balance) glued
-    text = re.sub(
-        r"(>[^<]*\bDESCRIPTION)\s*(Beginning Balance|Ending Balance)\b([^<]*<)",
-        r"\1\3",
-        text,
-        flags=re.IGNORECASE,
-    )
 
-    # BALANCE with a number glued or appended (keep the word BALANCE only)
-    text = re.sub(
-        r"(>[^<]*\bBALANCE)\s*\$?\s*\d{1,3}(?:,\d{3})*(?:\.\d{2})?\s*([^<]*<)",
-        r"\1\2",
-        text,
-        flags=re.IGNORECASE,
-    )
-    text = re.sub(
-        r"(>[^<]*\bBALANCE)\$",
-        r"\1",
-        text,
-        flags=re.IGNORECASE,
-    )
+def md_table_to_html(block: str) -> str:
+    lines = [ln.strip() for ln in block.strip().splitlines() if ln.strip()]
+    if len(lines) < 2:
+        return block
+
+    def split_row(row: str):
+        row = row.strip()
+        if row.startswith("|"):
+            row = row[1:]
+        if row.endswith("|"):
+            row = row[:-1]
+        return [p.strip() for p in row.split("|")]
+
+    header = split_row(lines[0])
+    body_lines = [ln for ln in lines[2:] if "|" in ln]
+
+    html_rows = []
+    html_rows.append("<tr>" + "".join(f"<th>{html.escape(c)}</th>" for c in header) + "</tr>")
+    for ln in body_lines:
+        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>")
+    return "<table>\n" + "\n".join(html_rows) + "\n</table>"
+
+
+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)
 
-    return text
+    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
 
 
-# ---- Colspan/rowspan expansion: align header and data rows for any PDF ----
+# ---- Generic HTML table normalizer (colspan/rowspan + beginning-balance repair) ----
 class TableGridParser(HTMLParser):
-    """Parse <table> HTML into rows of (text, colspan, rowspan)."""
+    """Parse a <table> into rows of (text, colspan, rowspan) per cell."""
 
     def __init__(self):
         super().__init__()
@@ -311,6 +306,7 @@ class TableGridParser(HTMLParser):
         self._cell_text = []
         self._colspan = 1
         self._rowspan = 1
+        self._in_cell = False
 
     def handle_starttag(self, tag, attrs):
         if tag == "tr":
@@ -320,20 +316,22 @@ class TableGridParser(HTMLParser):
             self._colspan = max(1, int(attrs_d.get("colspan", 1)))
             self._rowspan = max(1, int(attrs_d.get("rowspan", 1)))
             self._cell_text = []
+            self._in_cell = True
 
     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))
+            self._in_cell = False
         elif tag == "tr":
             self.rows.append(self._current_row)
 
     def handle_data(self, data):
-        self._cell_text.append(data)
+        if self._in_cell:
+            self._cell_text.append(data)
 
 
 def _build_grid(rows_data):
-    """Expand colspan/rowspan into a rectangular grid."""
     if not rows_data:
         return []
     blocked = defaultdict(set)
@@ -360,13 +358,12 @@ def _build_grid(rows_data):
 
 
 def _grid_to_html(grid):
-    """Emit a normalized table without colspan/rowspan."""
     if not grid:
         return ""
     lines = ["<table>"]
-    for i, row in enumerate(grid):
+    for r, row in enumerate(grid):
         lines.append("<tr>")
-        tag = "th" if i == 0 else "td"
+        tag = "th" if r == 0 else "td"
         for cell in row:
             escaped = (cell or "").replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;").replace('"', "&quot;")
             lines.append(f"<{tag}>{escaped}</{tag}>")
@@ -375,90 +372,107 @@ def _grid_to_html(grid):
     return "\n".join(lines)
 
 
-def expand_colspan_rowspan_in_tables(text):
-    """Normalize every <table>...</table> in text."""
-    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
+def _looks_like_amount(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
 
-    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 _normalize_transaction_header_and_beginning_balance(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.
+    """
+    if not grid or not grid[0]:
+        return grid
 
+    header = [c.strip() for c in grid[0]]
+    header_join = " | ".join(h.upper() for h in header)
 
-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()]
-    if len(lines) < 2:
-        return False
-    if "|" not in lines[0]:
-        return False
-    sep = lines[1].replace(" ", "")
-    return ("---" in sep) and ("|" in sep)
+    is_txn_like = ("DATE" in header_join and "AMOUNT" in header_join and "BALANCE" in header_join)
+    if not is_txn_like:
+        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
+    )
 
-def md_table_to_html(block: str) -> str:
-    """Convert a simple markdown pipe table to HTML table (best-effort)."""
-    lines = [ln.strip() for ln in block.strip().splitlines() if ln.strip()]
-    if len(lines) < 2:
-        return block
+    # 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)
 
-    def split_row(row: str):
-        row = row.strip()
-        if row.startswith("|"):
-            row = row[1:]
-        if row.endswith("|"):
-            row = row[:-1]
-        return [p.strip() for p in row.split("|")]
+    return grid
 
-    header = split_row(lines[0])
-    body_lines = [ln for ln in lines[2:] if "|" in ln]
 
-    html_rows = []
-    html_rows.append("<tr>" + "".join(f"<th>{html.escape(c)}</th>" for c in header) + "</tr>")
-    for ln in body_lines:
-        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>")
-    return "<table>\n" + "\n".join(html_rows) + "\n</table>"
+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
+      - emit normalized <table> HTML without colspan/rowspan
+    """
+    if not text or "<table" not in text.lower():
+        return text
 
+    pattern = re.compile(r"<table[^>]*>.*?</table>", re.DOTALL | re.IGNORECASE)
+    out = []
+    last = 0
 
-def normalize_money_glyphs(text: str) -> str:
-    """Conservative normalization for OCR number quirks."""
-    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)
+    for m in pattern.finditer(text):
+        out.append(text[last : m.start()])
+        table_html = m.group(0)
+        try:
+            p = TableGridParser()
+            p.feed(table_html)
+            grid = _build_grid(p.rows)
+            grid = _normalize_transaction_header_and_beginning_balance(grid)
+            out.append(_grid_to_html(grid) if grid else table_html)
+        except Exception:
+            out.append(table_html)
+        last = m.end()
 
-    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
+    out.append(text[last:])
+    return "".join(out)
 
 
 def stabilize_tables_and_text(page_md: str) -> str:
-    """Convert markdown pipe tables to HTML, clean header artifacts, normalize tables, and close tags."""
     if not page_md:
         return page_md
+
     page_md = normalize_money_glyphs(page_md)
 
     blocks = re.split(r"\n\s*\n", page_md.strip())
@@ -471,13 +485,10 @@ def stabilize_tables_and_text(page_md: str) -> str:
 
     stabilized = "\n\n".join(out_blocks)
 
-    # 1) Fix common header text artifacts
-    stabilized = clean_table_header_artifacts(stabilized)
-
-    # 2) Normalize tables to a rectangular grid (expand colspan/rowspan)
-    stabilized = expand_colspan_rowspan_in_tables(stabilized)
+    # Normalize all HTML tables (generic): expand colspan/rowspan + fix header artifacts + re-add beginning balance row when needed
+    stabilized = normalize_html_tables(stabilized)
 
-    # 3) Close any unclosed tags to prevent bleed
+    # Close tags to prevent bleed
     return close_unclosed_html(stabilized)
 
 
@@ -577,13 +588,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)
@@ -594,11 +603,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:
@@ -621,7 +628,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):