Update app.py

app.py

@@ -2,6 +2,7 @@
 
 import logging
 import os
+import re
 import tempfile
 import uuid
 from typing import List, Tuple
@@ -10,10 +11,9 @@ log = logging.getLogger("glmocr_simple_app")
 logging.basicConfig(level=logging.INFO)
 
 # ── Fine-tuned model repo on HuggingFace ─────────────────────────────────────
-# Loads from HF Hub at runtime — no local storage needed in the Space
 MERGED_MODEL_DIR = os.environ.get("MODEL_DIR", "SimpleCodeAI/glm-ocr-finetuned")
 
-
+RENDER_SCALE = 2.0
 PAD_LEFT_FRAC = 0.035
 PAD_RIGHT_FRAC = 0.10
 PAD_TOP_FRAC = 0.018
@@ -27,9 +27,8 @@ UNSHARP_PERCENT = 76
 UNSHARP_THRESHOLD = 1
 
 PAGE_PNG_COMPRESS_LEVEL = 3
-RENDER_SCALE = 2.0        # was 3.0 — reduce render size so resize isn't as aggressive
-MAX_IMAGE_SIDE = 1568     # was 1344 — allow slightly larger input to model
-MAX_NEW_TOKENS = 3000  # resize longest side to this before inference
+MAX_IMAGE_SIDE = 1568
+MAX_NEW_TOKENS = 3000
 
 # ── Model singleton ───────────────────────────────────────────────────────────
 _model = None
@@ -88,6 +87,98 @@ def _resize_for_inference(img):
     return img.resize(new_size, Image.LANCZOS)
 
 
+def _normalize_tables(text: str) -> str:
+    """
+    Convert plain-text table rows (no pipe boundaries) into proper
+    markdown table rows with | boundaries.
+
+    Handles rows like:
+        10/01  CCD DEBIT, SOME MERCHANT  654.00
+        10/01  DEBIT  2,500.00
+    and header rows like:
+        POSTING DATE  DESCRIPTION  AMOUNT
+    """
+    lines = text.split('\n')
+    result = []
+    i = 0
+
+    # Amount pattern: number like 1,234.56 or 123.45
+    amt_re = re.compile(r'[\d,]+\.\d{2}$')
+    # Date pattern at start of line: MM/DD
+    date_re = re.compile(r'^\d{2}/\d{2}\s')
+    # Header pattern
+    header_re = re.compile(
+        r'^(POSTING\s+DATE|DATE)\s+(DESCRIPTION|SERIAL\s+NO\.?|NO\.?\s+CHECKS)',
+        re.IGNORECASE
+    )
+
+    in_plain_table = False
+
+    while i < len(lines):
+        line = lines[i]
+        stripped = line.strip()
+
+        # Skip empty lines — reset table tracking
+        if not stripped:
+            in_plain_table = False
+            result.append(line)
+            i += 1
+            continue
+
+        # Already a proper markdown table row — pass through
+        if stripped.startswith('|'):
+            in_plain_table = False
+            result.append(line)
+            i += 1
+            continue
+
+        # Detect plain-text table header row
+        if header_re.match(stripped):
+            # Split by 2+ spaces
+            parts = re.split(r'\s{2,}', stripped)
+            parts = [p.strip() for p in parts if p.strip()]
+            if len(parts) >= 2:
+                result.append('| ' + ' | '.join(parts) + ' |')
+                result.append('| ' + ' | '.join(['---'] * len(parts)) + ' |')
+                in_plain_table = True
+                i += 1
+                continue
+
+        # Detect plain-text data row starting with date
+        if date_re.match(stripped):
+            # Split by 2+ spaces to separate columns
+            parts = re.split(r'\s{2,}', stripped)
+            parts = [p.strip() for p in parts if p.strip()]
+
+            # If only one part (description runs together), try splitting off amount
+            if len(parts) == 1 and amt_re.search(stripped):
+                # Split last amount token off
+                m = re.search(r'^(.*?)\s+([\d,]+\.\d{2})$', stripped)
+                if m:
+                    parts = [m.group(1).strip(), m.group(2).strip()]
+
+            if len(parts) >= 2:
+                result.append('| ' + ' | '.join(parts) + ' |')
+                in_plain_table = True
+                i += 1
+                continue
+
+        # Detect subtotal lines inside a plain table
+        if in_plain_table and re.match(r'^Subtotal:', stripped, re.IGNORECASE):
+            m = re.match(r'^(Subtotal:)\s+([\d,]+\.\d{2})', stripped, re.IGNORECASE)
+            if m:
+                result.append(f'| | **{m.group(1)}** | **{m.group(2)}** |')
+                i += 1
+                continue
+
+        # Not a table row
+        in_plain_table = False
+        result.append(line)
+        i += 1
+
+    return '\n'.join(result)
+
+
 def _infer_image(image_path: str) -> str:
     """Run fine-tuned model on a single image file and return markdown string."""
     import torch
@@ -127,13 +218,19 @@ def _infer_image(image_path: str) -> str:
                 **inputs,
                 max_new_tokens=MAX_NEW_TOKENS,
                 do_sample=False,
-                repetition_penalty=1.1,
+                repetition_penalty=1.15,
+                no_repeat_ngram_size=5,
             )
+
         result = processor.decode(
             ids[0][inputs["input_ids"].shape[1]:],
             skip_special_tokens=True,
         )
-        return result.strip()
+
+        # Normalize plain-text tables to proper markdown tables
+        result = _normalize_tables(result.strip())
+
+        return result
 
     finally:
         try:
@@ -247,7 +344,5 @@ def _create_gradio_demo():
 
 
 if __name__ == "__main__":
-    # Force model to load at startup
     _load_model()
-
     _create_gradio_demo().launch()