Update app.py

app.py

@@ -5,51 +5,89 @@ import os
 import re
 import tempfile
 import uuid
-from typing import List, Tuple
+from dataclasses import dataclass
+from typing import List, Optional, Tuple
 
-log = logging.getLogger("glmocr_simple_app")
+log = logging.getLogger("glmocr_improved_app")
 logging.basicConfig(level=logging.INFO)
 
-# ── Fine-tuned model repo on HuggingFace ─────────────────────────────────────
+# Fine-tuned model repo on HuggingFace
 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
-PAD_BOTTOM_FRAC = 0.018
+# Primary render profile (fast)
+RENDER_SCALE = 2.6
+PAD_LEFT_FRAC = 0.04
+PAD_RIGHT_FRAC = 0.16  # right-side amounts are often clipped
+PAD_TOP_FRAC = 0.02
+PAD_BOTTOM_FRAC = 0.02
+
+# Retry render profile (quality)
+RETRY_RENDER_SCALE = 3.1
+RETRY_PAD_RIGHT_FRAC = 0.24
 
 ENABLE_CONTRAST = True
-CONTRAST_FACTOR = 1.18
+CONTRAST_FACTOR = 1.20
 ENABLE_UNSHARP = True
 UNSHARP_RADIUS = 0.78
-UNSHARP_PERCENT = 76
+UNSHARP_PERCENT = 82
 UNSHARP_THRESHOLD = 1
 
 PAGE_PNG_COMPRESS_LEVEL = 3
-MAX_IMAGE_SIDE = 1568
-MAX_NEW_TOKENS = 3000
+MAX_IMAGE_SIDE = 1792
+MAX_NEW_TOKENS = 3800
 
-# ── Model singleton ───────────────────────────────────────────────────────────
+# Model singleton
 _model = None
 _processor = None
 
 
+@dataclass
+class RenderProfile:
+    scale: float
+    pad_left: float
+    pad_right: float
+    pad_top: float
+    pad_bottom: float
+
+
+PRIMARY_PROFILE = RenderProfile(
+    scale=RENDER_SCALE,
+    pad_left=PAD_LEFT_FRAC,
+    pad_right=PAD_RIGHT_FRAC,
+    pad_top=PAD_TOP_FRAC,
+    pad_bottom=PAD_BOTTOM_FRAC,
+)
+RETRY_PROFILE = RenderProfile(
+    scale=RETRY_RENDER_SCALE,
+    pad_left=PAD_LEFT_FRAC,
+    pad_right=RETRY_PAD_RIGHT_FRAC,
+    pad_top=PAD_TOP_FRAC,
+    pad_bottom=PAD_BOTTOM_FRAC,
+)
+
+
 def _load_model():
     global _model, _processor
     if _model is not None:
         return _model, _processor
 
     import torch
-    from transformers import AutoProcessor, AutoModelForImageTextToText
+    from transformers import AutoModelForImageTextToText, AutoProcessor
 
     log.info("Loading fine-tuned model from %s ...", MERGED_MODEL_DIR)
     _processor = AutoProcessor.from_pretrained(
-        MERGED_MODEL_DIR, trust_remote_code=True
+        MERGED_MODEL_DIR,
+        trust_remote_code=True,
     )
+
+    dtype = torch.bfloat16
+    if not torch.cuda.is_available():
+        # CPU path is safer with fp32
+        dtype = torch.float32
+
     _model = AutoModelForImageTextToText.from_pretrained(
         MERGED_MODEL_DIR,
-        dtype=torch.bfloat16,
+        dtype=dtype,
         device_map="auto",
         trust_remote_code=True,
     )
@@ -89,98 +127,138 @@ def _resize_for_inference(img):
 
 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
+    Normalize plain-text table rows into markdown table rows.
+    This is critical for downstream parsers that expect markdown-table structure.
     """
-    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
+    lines = text.split("\n")
+    result: List[str] = []
+
+    amt_re = re.compile(r"[\d,]+\.\d{2}$")
+    date_re = re.compile(r"^\d{2}/\d{2}\s")
     header_re = re.compile(
-        r'^(POSTING\s+DATE|DATE)\s+(DESCRIPTION|SERIAL\s+NO\.?|NO\.?\s+CHECKS)',
-        re.IGNORECASE
+        r"^(POSTING\s+DATE|DATE)\s+(DESCRIPTION|SERIAL\s+NO\.?|NO\.?\s+CHECKS)",
+        re.IGNORECASE,
     )
+    subtotal_re = re.compile(r"^(Subtotal:)\s+([\d,]+\.\d{2})", re.IGNORECASE)
 
     in_plain_table = False
 
-    while i < len(lines):
-        line = lines[i]
+    for raw_line in lines:
+        line = raw_line.rstrip()
         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('|'):
+        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()]
+            parts = [p.strip() for p in re.split(r"\s{2,}", stripped) if p.strip()]
             if len(parts) >= 2:
-                result.append('| ' + ' | '.join(parts) + ' |')
-                result.append('| ' + ' | '.join(['---'] * len(parts)) + ' |')
+                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()]
+            parts = [p.strip() for p in re.split(r"\s{2,}", stripped) if p.strip()]
 
-            # If only one part (description runs together), try splitting off amount
+            # Recover lines where spacing collapsed and only amount is clearly separable.
             if len(parts) == 1 and amt_re.search(stripped):
-                # Split last amount token off
-                m = re.search(r'^(.*?)\s+([\d,]+\.\d{2})$', stripped)
+                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) + ' |')
+                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 in_plain_table:
+            m = subtotal_re.match(stripped)
             if m:
-                result.append(f'| | **{m.group(1)}** | **{m.group(2)}** |')
-                i += 1
+                result.append(f"| | **{m.group(1)}** | **{m.group(2)}** |")
                 continue
 
-        # Not a table row
         in_plain_table = False
         result.append(line)
-        i += 1
 
-    return '\n'.join(result)
+    return "\n".join(result)
+
+
+def _strip_code_fences(text: str) -> str:
+    # Some generations wrap markdown in code fences; strip only outer wrappers.
+    t = text.strip()
+    if t.startswith("```") and t.endswith("```"):
+        t = re.sub(r"^```[a-zA-Z0-9_-]*\n?", "", t)
+        t = re.sub(r"\n?```$", "", t)
+    return t.strip()
 
 
-def _infer_image(image_path: str) -> str:
-    """Run fine-tuned model on a single image file and return markdown string."""
+def _postprocess_markdown(text: str) -> str:
+    text = _strip_code_fences(text)
+    text = _normalize_tables(text)
+    return text.strip()
+
+
+def _build_prompt(strict_table_mode: bool = False) -> str:
+    base = (
+        "Document Parsing to markdown.\n"
+        "Rules:\n"
+        "1) Preserve every row and amount exactly as seen.\n"
+        "2) Keep transaction/payment/check sections in markdown table format.\n"
+        "3) Never drop right-most numeric columns (amount/balance).\n"
+        "4) Keep page content order exactly.\n"
+        "5) Do not summarize.\n"
+    )
+    if strict_table_mode:
+        base += (
+            "6) If a row starts with a date (MM/DD), output it as a table row and include trailing amount.\n"
+            "7) Prefer explicit table rows over plain text for statement activity.\n"
+        )
+    return base
+
+
+def _count_amounts(text: str) -> int:
+    return len(re.findall(r"\b\d{1,3}(?:,\d{3})*\.\d{2}\b", text))
+
+
+def _count_dated_rows(text: str) -> int:
+    return len(re.findall(r"(?m)^\s*\d{2}/\d{2}\b", text))
+
+
+def _looks_amount_missing(text: str) -> bool:
+    """
+    Heuristic: many dated activity rows but very low amount density.
+    This catches pages where payment rows were parsed but right-side amounts vanished.
+    """
+    low = text.lower()
+    likely_activity = any(
+        key in low
+        for key in (
+            "daily account activity",
+            "electronic payments",
+            "electronic deposits",
+            "checks paid",
+            "other withdrawals",
+            "service charges",
+        )
+    )
+    if not likely_activity:
+        return False
+
+    dated = _count_dated_rows(text)
+    amts = _count_amounts(text)
+    return dated >= 12 and amts <= max(3, dated // 10)
+
+
+def _infer_image(image_path: str, strict_table_mode: bool = False) -> str:
+    """Run fine-tuned model on a single image file and return markdown."""
     import torch
     from PIL import Image
 
@@ -194,13 +272,15 @@ def _infer_image(image_path: str) -> str:
     try:
         img.save(resized_path, "PNG")
 
-        messages = [{
-            "role": "user",
-            "content": [
-                {"type": "image", "url": resized_path},
-                {"type": "text", "text": "Document Parsing:"},
-            ],
-        }]
+        messages = [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "image", "url": resized_path},
+                    {"type": "text", "text": _build_prompt(strict_table_mode)},
+                ],
+            }
+        ]
 
         inputs = processor.apply_chat_template(
             messages,
@@ -211,25 +291,22 @@ def _infer_image(image_path: str) -> str:
         ).to(model.device)
         inputs.pop("token_type_ids", None)
 
-        torch.cuda.empty_cache()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
 
         with torch.no_grad():
             ids = model.generate(
                 **inputs,
                 max_new_tokens=MAX_NEW_TOKENS,
                 do_sample=False,
-                repetition_penalty=1.1,
+                repetition_penalty=1.08,
             )
 
         result = processor.decode(
-            ids[0][inputs["input_ids"].shape[1]:],
+            ids[0][inputs["input_ids"].shape[1] :],
             skip_special_tokens=True,
         )
-
-        result = result.strip()
-
-        return result
-
+        return _postprocess_markdown(result)
     finally:
         try:
             os.unlink(resized_path)
@@ -237,69 +314,95 @@ def _infer_image(image_path: str) -> str:
             pass
 
 
-def render_pdf_pages_to_images(pdf_path: str) -> Tuple[List[str], List[int]]:
+def _render_pdf_page_to_image(pdf_path: str, page_index: int, profile: RenderProfile) -> str:
     import pymupdf as fitz
     from PIL import Image
 
     doc = fitz.open(pdf_path)
-    page_images: List[str] = []
-    page_heights: List[int] = []
+    try:
+        page = doc[page_index]
+        pix = page.get_pixmap(matrix=fitz.Matrix(profile.scale, profile.scale), alpha=False)
+    finally:
+        doc.close()
 
-    for i in range(len(doc)):
-        page = doc[i]
-        pix = page.get_pixmap(
-            matrix=fitz.Matrix(RENDER_SCALE, RENDER_SCALE), alpha=False
-        )
+    img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
+    img = _enhance_raster_for_ocr(img)
 
-        img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
-        img = _enhance_raster_for_ocr(img)
+    w, h = img.size
+    pad_l = int(w * profile.pad_left)
+    pad_r = int(w * profile.pad_right)
+    pad_t = int(h * profile.pad_top)
+    pad_b = int(h * profile.pad_bottom)
+
+    if any(p > 0 for p in (pad_l, pad_r, pad_t, pad_b)):
+        canvas = Image.new(
+            "RGB",
+            (w + pad_l + pad_r, h + pad_t + pad_b),
+            (255, 255, 255),
+        )
+        canvas.paste(img, (pad_l, pad_t))
+        img = canvas
 
-        w, h = img.size
-        pad_l = int(w * PAD_LEFT_FRAC)
-        pad_r = int(w * PAD_RIGHT_FRAC)
-        pad_t = int(h * PAD_TOP_FRAC)
-        pad_b = int(h * PAD_BOTTOM_FRAC)
+    uniq = uuid.uuid4().hex[:10]
+    out_path = os.path.join(
+        tempfile.gettempdir(),
+        f"glmocr_page_{os.getpid()}_{uniq}_{page_index}.png",
+    )
+    img.save(out_path, "PNG", compress_level=PAGE_PNG_COMPRESS_LEVEL)
+    return out_path
 
-        if any(p > 0 for p in (pad_l, pad_r, pad_t, pad_b)):
-            canvas = Image.new(
-                "RGB", (w + pad_l + pad_r, h + pad_t + pad_b), (255, 255, 255)
-            )
-            canvas.paste(img, (pad_l, pad_t))
-            img = canvas
 
-        uniq = uuid.uuid4().hex[:10]
-        img_path = os.path.join(
-            tempfile.gettempdir(),
-            f"glmocr_page_{os.getpid()}_{uniq}_{i}.png",
-        )
-        img.save(img_path, "PNG", compress_level=PAGE_PNG_COMPRESS_LEVEL)
-        page_images.append(img_path)
-        page_heights.append(img.height)
+def _pdf_page_count(pdf_path: str) -> int:
+    import pymupdf as fitz
 
-    doc.close()
-    return page_images, page_heights
+    doc = fitz.open(pdf_path)
+    try:
+        return len(doc)
+    finally:
+        doc.close()
 
 
 def run_ocr(uploaded_file):
     if uploaded_file is None:
         return "Please upload a file."
 
-    page_images: List[str] = []
+    temp_paths: List[str] = []
     try:
         path = uploaded_file.name if hasattr(uploaded_file, "name") else str(uploaded_file)
         is_pdf = path.lower().endswith(".pdf")
 
-        if is_pdf:
-            page_images, _ = render_pdf_pages_to_images(path)
-        else:
-            page_images = [path]
+        all_pages: List[str] = []
 
-        all_pages = []
-        for page_num, img_path in enumerate(page_images):
-            log.info("Processing page %d / %d ...", page_num + 1, len(page_images))
-            page_md = _infer_image(img_path)
+        if not is_pdf:
+            page_md = _infer_image(path, strict_table_mode=True)
             if page_md:
                 all_pages.append(page_md)
+        else:
+            page_total = _pdf_page_count(path)
+            for page_idx in range(page_total):
+                log.info("Processing page %d / %d ...", page_idx + 1, page_total)
+
+                # Pass 1: primary profile
+                img_path = _render_pdf_page_to_image(path, page_idx, PRIMARY_PROFILE)
+                temp_paths.append(img_path)
+                page_md = _infer_image(img_path, strict_table_mode=False)
+
+                # Pass 2: retry for amount-missing pages
+                if _looks_amount_missing(page_md):
+                    log.info(
+                        "Page %d flagged as amount-missing; retrying high-quality render.",
+                        page_idx + 1,
+                    )
+                    retry_img = _render_pdf_page_to_image(path, page_idx, RETRY_PROFILE)
+                    temp_paths.append(retry_img)
+                    retry_md = _infer_image(retry_img, strict_table_mode=True)
+
+                    # Keep better output by amount density.
+                    if _count_amounts(retry_md) > _count_amounts(page_md):
+                        page_md = retry_md
+
+                if page_md:
+                    all_pages.append(page_md)
 
         merged = (
             "\n\n---page-separator---\n\n".join(all_pages)
@@ -310,11 +413,12 @@ 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()}"
 
     finally:
-        for p in page_images:
+        for p in temp_paths:
             try:
                 if (
                     isinstance(p, str)
@@ -329,8 +433,8 @@ def run_ocr(uploaded_file):
 def _create_gradio_demo():
     import gradio as gr
 
-    with gr.Blocks(title="GLM-OCR Fine-tuned") as demo:
-        gr.Markdown("# GLM-OCR (Fine-tuned)")
+    with gr.Blocks(title="GLM-OCR Fine-tuned (Improved)") as demo:
+        gr.Markdown("# GLM-OCR (Fine-tuned, Improved)")
         file_in = gr.File(
             label="Upload PDF or image",
             file_types=[".pdf", ".png", ".jpg", ".jpeg", ".tiff", ".bmp"],
@@ -343,4 +447,4 @@ def _create_gradio_demo():
 
 if __name__ == "__main__":
     _load_model()
-    _create_gradio_demo().launch()
+    _create_gradio_demo().launch()