Update app.py

app.py

@@ -2,7 +2,6 @@
 
 import logging
 import os
-import re
 import tempfile
 import uuid
 from typing import List, Tuple
@@ -10,7 +9,7 @@ from typing import List, Tuple
 log = logging.getLogger("glmocr_simple_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
@@ -19,10 +18,6 @@ PAD_RIGHT_FRAC = 0.10
 PAD_TOP_FRAC = 0.018
 PAD_BOTTOM_FRAC = 0.018
 
-# 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.18
 ENABLE_UNSHARP = True
@@ -34,7 +29,7 @@ PAGE_PNG_COMPRESS_LEVEL = 3
 MAX_IMAGE_SIDE = 1568
 MAX_NEW_TOKENS = 3000
 
-# Model singleton
+# ── Model singleton ───────────────────────────────────────────────────────────
 _model = None
 _processor = None
 
@@ -45,16 +40,15 @@ def _load_model():
         return _model, _processor
 
     import torch
-    from transformers import AutoModelForImageTextToText, AutoProcessor
+    from transformers import AutoProcessor, AutoModelForImageTextToText
 
     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
     )
     _model = AutoModelForImageTextToText.from_pretrained(
         MERGED_MODEL_DIR,
-        dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
+        dtype=torch.bfloat16,
         device_map="auto",
         trust_remote_code=True,
     )
@@ -92,69 +86,8 @@ def _resize_for_inference(img):
     return img.resize(new_size, Image.LANCZOS)
 
 
-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)
-        t = re.sub(r"\n?```$", "", t)
-    return t.strip()
-
-
-def _build_prompt(strict: bool = False) -> str:
-    base = (
-        "Document Parsing to markdown.\n"
-        "Rules:\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:
-        base += (
-            "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 _is_hallucinated(text: str) -> bool:
-    low = text.lower()
-    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:
-    low = text.lower()
-    if not any(
-        key in low
-        for key in (
-            "electronic payments",
-            "electronic deposits",
-            "daily account activity",
-            "checks paid",
-            "other withdrawals",
-        )
-    ):
-        return False
-    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: bool = False) -> str:
-    """Run model on a single image and return markdown string."""
+def _infer_image(image_path: str) -> str:
+    """Run fine-tuned model on a single image file and return markdown string."""
     import torch
     from PIL import Image
 
@@ -168,15 +101,13 @@ def _infer_image(image_path: str, strict: bool = False) -> str:
     try:
         img.save(resized_path, "PNG")
 
-        messages = [
-            {
-                "role": "user",
-                "content": [
-                    {"type": "image", "url": resized_path},
-                    {"type": "text", "text": _build_prompt(strict)},
-                ],
-            }
-        ]
+        messages = [{
+            "role": "user",
+            "content": [
+                {"type": "image", "url": resized_path},
+                {"type": "text", "text": "Document Parsing:"},
+            ],
+        }]
 
         inputs = processor.apply_chat_template(
             messages,
@@ -187,8 +118,7 @@ def _infer_image(image_path: str, strict: bool = False) -> str:
         ).to(model.device)
         inputs.pop("token_type_ids", None)
 
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
+        torch.cuda.empty_cache()
 
         with torch.no_grad():
             ids = model.generate(
@@ -199,10 +129,10 @@ def _infer_image(image_path: str, strict: bool = False) -> str:
             )
 
         result = processor.decode(
-            ids[0][inputs["input_ids"].shape[1] :],
+            ids[0][inputs["input_ids"].shape[1]:],
             skip_special_tokens=True,
         )
-        return _strip_outer_code_fences(result)
+        return result.strip()
 
     finally:
         try:
@@ -211,55 +141,46 @@ def _infer_image(image_path: str, strict: bool = False) -> str:
             pass
 
 
-def _render_page(pdf_path: str, page_idx: int, scale: float, right_pad_frac: float) -> str:
+def render_pdf_pages_to_images(pdf_path: str) -> Tuple[List[str], List[int]]:
     import pymupdf as fitz
     from PIL import Image
 
     doc = fitz.open(pdf_path)
-    try:
-        page = doc[page_idx]
-        pix = page.get_pixmap(matrix=fitz.Matrix(scale, scale), alpha=False)
-    finally:
-        doc.close()
-
-    img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
-    img = _enhance_raster_for_ocr(img)
-
-    w, h = img.size
-    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.paste(img, (pad_l, pad_t))
-        img = canvas
+    page_images: List[str] = []
+    page_heights: List[int] = []
 
-    uniq = uuid.uuid4().hex[:10]
-    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
+    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)
 
-def _pdf_page_count(pdf_path: str) -> int:
-    import pymupdf as fitz
+        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)
 
-    doc = fitz.open(pdf_path)
-    try:
-        return len(doc)
-    finally:
-        doc.close()
+        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 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)
+    doc.close()
     return page_images, page_heights
 
 
@@ -273,54 +194,26 @@ def run_ocr(uploaded_file):
         is_pdf = path.lower().endswith(".pdf")
 
         if is_pdf:
-            page_total = _pdf_page_count(path)
-            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
-
-        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
+            page_images, _ = render_pdf_pages_to_images(path)
+        else:
+            page_images = [path]
+
+        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 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
 
     except Exception as e:
         import traceback
-
         log.exception("run_ocr failed: %s", e)
         return f"Error: {e}\n\n{traceback.format_exc()}"
 
@@ -354,4 +247,4 @@ def _create_gradio_demo():
 
 if __name__ == "__main__":
     _load_model()
-    _create_gradio_demo().launch()
+    _create_gradio_demo().launch()