Update app.py

app.py

@@ -5,67 +5,40 @@ import os
 import re
 import tempfile
 import uuid
-from dataclasses import dataclass
-from typing import List, Optional, Tuple
+from typing import List, Tuple
 
-log = logging.getLogger("glmocr_improved_app")
+log = logging.getLogger("glmocr_simple_app")
 logging.basicConfig(level=logging.INFO)
 
 # Fine-tuned model repo on HuggingFace
 MERGED_MODEL_DIR = os.environ.get("MODEL_DIR", "SimpleCodeAI/glm-ocr-finetuned")
 
-# 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
+RENDER_SCALE = 2.0
+PAD_LEFT_FRAC = 0.035
+PAD_RIGHT_FRAC = 0.10
+PAD_TOP_FRAC = 0.018
+PAD_BOTTOM_FRAC = 0.018
 
-# Retry render profile (quality)
-RETRY_RENDER_SCALE = 3.1
-RETRY_PAD_RIGHT_FRAC = 0.24
+# Slightly stronger retry profile for right-side amount clipping
+RETRY_RENDER_SCALE = 2.6
+RETRY_PAD_RIGHT_FRAC = 0.18
 
 ENABLE_CONTRAST = True
-CONTRAST_FACTOR = 1.20
+CONTRAST_FACTOR = 1.18
 ENABLE_UNSHARP = True
 UNSHARP_RADIUS = 0.78
-UNSHARP_PERCENT = 82
+UNSHARP_PERCENT = 76
 UNSHARP_THRESHOLD = 1
 
 PAGE_PNG_COMPRESS_LEVEL = 3
-MAX_IMAGE_SIDE = 1792
-MAX_NEW_TOKENS = 3800
+MAX_IMAGE_SIDE = 1568
+MAX_NEW_TOKENS = 3000
 
 # 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:
@@ -79,15 +52,9 @@ def _load_model():
         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=dtype,
+        dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
         device_map="auto",
         trust_remote_code=True,
     )
@@ -125,74 +92,7 @@ def _resize_for_inference(img):
     return img.resize(new_size, Image.LANCZOS)
 
 
-def _normalize_tables(text: str) -> str:
-    """
-    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: 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,
-    )
-    subtotal_re = re.compile(r"^(Subtotal:)\s+([\d,]+\.\d{2})", re.IGNORECASE)
-
-    in_plain_table = False
-
-    for raw_line in lines:
-        line = raw_line.rstrip()
-        stripped = line.strip()
-
-        if not stripped:
-            in_plain_table = False
-            result.append(line)
-            continue
-
-        if stripped.startswith("|"):
-            in_plain_table = False
-            result.append(line)
-            continue
-
-        if header_re.match(stripped):
-            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)) + " |")
-                in_plain_table = True
-                continue
-
-        if date_re.match(stripped):
-            parts = [p.strip() for p in re.split(r"\s{2,}", stripped) if p.strip()]
-
-            # Recover lines where spacing collapsed and only amount is clearly separable.
-            if len(parts) == 1 and amt_re.search(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) + " |")
-                in_plain_table = True
-                continue
-
-        if in_plain_table:
-            m = subtotal_re.match(stripped)
-            if m:
-                result.append(f"| | **{m.group(1)}** | **{m.group(2)}** |")
-                continue
-
-        in_plain_table = False
-        result.append(line)
-
-    return "\n".join(result)
-
-
-def _strip_code_fences(text: str) -> str:
-    # Some generations wrap markdown in code fences; strip only outer wrappers.
+def _strip_outer_code_fences(text: str) -> str:
     t = text.strip()
     if t.startswith("```") and t.endswith("```"):
         t = re.sub(r"^```[a-zA-Z0-9_-]*\n?", "", t)
@@ -200,136 +100,61 @@ def _strip_code_fences(text: str) -> str:
     return t.strip()
 
 
-def _postprocess_markdown(text: str) -> str:
-    text = _strip_code_fences(text)
-    text = _html_to_markdown_tables(text)
-    text = _normalize_tables(text)
-    return text.strip()
-
-
-def _build_prompt(strict_table_mode: bool = False) -> str:
+def _build_prompt(strict: 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"
-        "6) Output markdown only.\n"
-        "7) Do not output HTML tags (<table>, <tr>, <td>, <thead>, <tbody>, etc.).\n"
-        "8) Do not invent sample/demo rows.\n"
+        "1) Preserve content exactly in reading order.\n"
+        "2) Do not summarize.\n"
+        "3) Do not invent rows, values, subtotals, or examples.\n"
+        "4) Output markdown only; do not output HTML tags.\n"
     )
-    if strict_table_mode:
+    if strict:
         base += (
-            "9) If a row starts with a date (MM/DD), output it as a table row and include trailing amount.\n"
-            "10) Prefer explicit table rows over plain text for statement activity.\n"
+            "5) For statement activity tables, preserve date/description/amount rows exactly.\n"
+            "6) If uncertain, keep raw line text instead of fabricating table rows.\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 _html_to_markdown_tables(text: str) -> str:
-    """
-    Convert basic HTML table blocks into pipe markdown tables.
-    This keeps downstream parsing deterministic even if model emits HTML.
-    """
-    table_re = re.compile(r"<table[\s\S]*?</table>", re.IGNORECASE)
-    row_re = re.compile(r"<tr[\s\S]*?</tr>", re.IGNORECASE)
-    cell_re = re.compile(r"<t[dh][^>]*>([\s\S]*?)</t[dh]>", re.IGNORECASE)
-    br_re = re.compile(r"<br\s*/?>", re.IGNORECASE)
-    tag_re = re.compile(r"<[^>]+>")
-
-    def _clean_cell(s: str) -> str:
-        s = br_re.sub(" ", s)
-        s = tag_re.sub(" ", s)
-        s = s.replace("&nbsp;", " ")
-        s = s.replace("&amp;", "&")
-        s = re.sub(r"\s+", " ", s).strip()
-        return s
-
-    def _convert_one(table_html: str) -> str:
-        rows = []
-        for row_html in row_re.findall(table_html):
-            cells = [_clean_cell(c) for c in cell_re.findall(row_html)]
-            cells = [c for c in cells if c != ""]
-            if cells:
-                rows.append(cells)
-        if not rows:
-            return table_html
-
-        width = max(len(r) for r in rows)
-        norm_rows = [r + [""] * (width - len(r)) for r in rows]
-        header = norm_rows[0]
-        out = [
-            "| " + " | ".join(header) + " |",
-            "| " + " | ".join(["---"] * width) + " |",
-        ]
-        for r in norm_rows[1:]:
-            out.append("| " + " | ".join(r) + " |")
-        return "\n".join(out)
-
-    return table_re.sub(lambda m: _convert_one(m.group(0)), text)
-
-
-def _is_hallucinated_table_block(text: str) -> bool:
-    """
-    Detect obvious synthetic outputs like repeated '100.00' rows or
-    placeholder transaction lists not present in statement pages.
-    """
+def _is_hallucinated(text: str) -> bool:
     low = text.lower()
-    if "<table" not in low and "| transaction | amount |" not in low:
-        return False
-
-    repeated_100 = len(re.findall(r"\b100\.00\b", text))
-    credit_card_rows = len(re.findall(r"\bcredit card\b", low))
-    bank_transfer_rows = len(re.findall(r"\bbank transfer\b", low))
-
-    if repeated_100 >= 18 and (credit_card_rows + bank_transfer_rows) >= 10:
+    if "<table" in low or "<thead" in low or "<tbody" in low:
         return True
     if re.search(r"<!--\s*row\s+\d+\s*-->", low):
         return True
+    # Signature from your bad sample
+    if (
+        len(re.findall(r"\b100\.00\b", text)) >= 15
+        and len(re.findall(r"\bcredit card\b|\bbank transfer\b", low)) >= 8
+    ):
+        return True
     return False
 
 
+def _count_amounts(text: str) -> int:
+    return len(re.findall(r"\b\d{1,3}(?:,\d{3})*\.\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(
+    if not any(
         key in low
         for key in (
-            "daily account activity",
             "electronic payments",
             "electronic deposits",
+            "daily account activity",
             "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)
+    dated_rows = len(re.findall(r"(?m)^\s*\d{2}/\d{2}\b", text))
+    return dated_rows >= 10 and _count_amounts(text) <= max(2, dated_rows // 12)
 
 
-def _infer_image(
-    image_path: str,
-    strict_table_mode: bool = False,
-    conservative_mode: bool = False,
-) -> str:
-    """Run fine-tuned model on a single image file and return markdown."""
+def _infer_image(image_path: str, strict: bool = False) -> str:
+    """Run model on a single image and return markdown string."""
     import torch
     from PIL import Image
 
@@ -343,19 +168,12 @@ def _infer_image(
     try:
         img.save(resized_path, "PNG")
 
-        prompt = _build_prompt(strict_table_mode)
-        if conservative_mode:
-            prompt += (
-                "\n11) If unsure, keep original text line-by-line; do not fabricate values.\n"
-                "12) Never generate template/example/sample transaction rows.\n"
-            )
-
         messages = [
             {
                 "role": "user",
                 "content": [
                     {"type": "image", "url": resized_path},
-                    {"type": "text", "text": prompt},
+                    {"type": "text", "text": _build_prompt(strict)},
                 ],
             }
         ]
@@ -377,14 +195,15 @@ def _infer_image(
                 **inputs,
                 max_new_tokens=MAX_NEW_TOKENS,
                 do_sample=False,
-                repetition_penalty=1.08,
+                repetition_penalty=1.1,
             )
 
         result = processor.decode(
             ids[0][inputs["input_ids"].shape[1] :],
             skip_special_tokens=True,
         )
-        return _postprocess_markdown(result)
+        return _strip_outer_code_fences(result)
+
     finally:
         try:
             os.unlink(resized_path)
@@ -392,14 +211,14 @@ def _infer_image(
             pass
 
 
-def _render_pdf_page_to_image(pdf_path: str, page_index: int, profile: RenderProfile) -> str:
+def _render_page(pdf_path: str, page_idx: int, scale: float, right_pad_frac: float) -> str:
     import pymupdf as fitz
     from PIL import Image
 
     doc = fitz.open(pdf_path)
     try:
-        page = doc[page_index]
-        pix = page.get_pixmap(matrix=fitz.Matrix(profile.scale, profile.scale), alpha=False)
+        page = doc[page_idx]
+        pix = page.get_pixmap(matrix=fitz.Matrix(scale, scale), alpha=False)
     finally:
         doc.close()
 
@@ -407,27 +226,20 @@ def _render_pdf_page_to_image(pdf_path: str, page_index: int, profile: RenderPro
     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)
+    pad_l = int(w * PAD_LEFT_FRAC)
+    pad_r = int(w * right_pad_frac)
+    pad_t = int(h * PAD_TOP_FRAC)
+    pad_b = int(h * PAD_BOTTOM_FRAC)
 
     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 = 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]
-    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
+    img_path = os.path.join(tempfile.gettempdir(), f"glmocr_page_{os.getpid()}_{uniq}_{page_idx}.png")
+    img.save(img_path, "PNG", compress_level=PAGE_PNG_COMPRESS_LEVEL)
+    return img_path
 
 
 def _pdf_page_count(pdf_path: str) -> int:
@@ -440,74 +252,71 @@ def _pdf_page_count(pdf_path: str) -> int:
         doc.close()
 
 
+def render_pdf_pages_to_images(pdf_path: str) -> Tuple[List[str], List[int]]:
+    page_images: List[str] = []
+    page_heights: List[int] = []
+    total = _pdf_page_count(pdf_path)
+    for i in range(total):
+        path = _render_page(pdf_path, i, RENDER_SCALE, PAD_RIGHT_FRAC)
+        page_images.append(path)
+        page_heights.append(0)
+    return page_images, page_heights
+
+
 def run_ocr(uploaded_file):
     if uploaded_file is None:
         return "Please upload a file."
 
-    temp_paths: List[str] = []
+    page_images: List[str] = []
     try:
         path = uploaded_file.name if hasattr(uploaded_file, "name") else str(uploaded_file)
         is_pdf = path.lower().endswith(".pdf")
 
-        all_pages: List[str] = []
-
-        if not is_pdf:
-            page_md = _infer_image(path, strict_table_mode=True, conservative_mode=True)
-            if _is_hallucinated_table_block(page_md):
-                page_md = _infer_image(path, strict_table_mode=True, conservative_mode=True)
-            if page_md:
-                all_pages.append(page_md)
-        else:
+        if is_pdf:
             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 1.5: recover from synthetic/hallucinated table outputs
-                if _is_hallucinated_table_block(page_md):
-                    log.info(
-                        "Page %d flagged as hallucinated-table; retrying conservative prompt.",
-                        page_idx + 1,
-                    )
-                    page_md_retry = _infer_image(
-                        img_path,
-                        strict_table_mode=True,
-                        conservative_mode=True,
-                    )
-                    if not _is_hallucinated_table_block(page_md_retry):
-                        page_md = page_md_retry
-
-                # 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,
-                        conservative_mode=True,
-                    )
-
-                    # Keep better output by amount density.
-                    if _count_amounts(retry_md) > _count_amounts(page_md):
+            all_pages: List[str] = []
+
+            for page_num in range(page_total):
+                log.info("Processing page %d / %d ...", page_num + 1, page_total)
+                img_path = _render_page(path, page_num, RENDER_SCALE, PAD_RIGHT_FRAC)
+                page_images.append(img_path)
+
+                page_md = _infer_image(img_path, strict=False)
+
+                # Retry when output looks clearly wrong.
+                if _is_hallucinated(page_md) or _looks_amount_missing(page_md):
+                    retry_img = _render_page(path, page_num, RETRY_RENDER_SCALE, RETRY_PAD_RIGHT_FRAC)
+                    page_images.append(retry_img)
+                    retry_md = _infer_image(retry_img, strict=True)
+
+                    # Prefer non-hallucinated output; otherwise keep original.
+                    if not _is_hallucinated(retry_md):
                         page_md = retry_md
 
+                # Hard guard: if still hallucinated, drop page instead of poisoning full markdown.
+                if _is_hallucinated(page_md):
+                    log.warning("Skipping hallucinated output on page %d", page_num + 1)
+                    continue
+
                 if page_md:
                     all_pages.append(page_md)
 
-        merged = (
-            "\n\n---page-separator---\n\n".join(all_pages)
-            if all_pages
-            else "(No content extracted)"
-        )
-        return merged
+            merged = (
+                "\n\n---page-separator---\n\n".join(all_pages)
+                if all_pages
+                else "(No content extracted)"
+            )
+            return merged
+
+        page_md = _infer_image(path, strict=True)
+        if _is_hallucinated(page_md):
+            # Last fallback for single image
+            page_md_retry = _infer_image(path, strict=True)
+            if not _is_hallucinated(page_md_retry):
+                page_md = page_md_retry
+        if _is_hallucinated(page_md):
+            return "(Output rejected: detected hallucinated HTML/sample table content)"
+        return page_md
 
     except Exception as e:
         import traceback
@@ -516,7 +325,7 @@ def run_ocr(uploaded_file):
         return f"Error: {e}\n\n{traceback.format_exc()}"
 
     finally:
-        for p in temp_paths:
+        for p in page_images:
             try:
                 if (
                     isinstance(p, str)
@@ -531,8 +340,8 @@ def run_ocr(uploaded_file):
 def _create_gradio_demo():
     import gradio as gr
 
-    with gr.Blocks(title="GLM-OCR Fine-tuned (Improved)") as demo:
-        gr.Markdown("# GLM-OCR (Fine-tuned, Improved)")
+    with gr.Blocks(title="GLM-OCR Fine-tuned") as demo:
+        gr.Markdown("# GLM-OCR (Fine-tuned)")
         file_in = gr.File(
             label="Upload PDF or image",
             file_types=[".pdf", ".png", ".jpg", ".jpeg", ".tiff", ".bmp"],