Update app.py

app.py

@@ -1,361 +1,87 @@
 #!/usr/bin/env python3
 
+import asyncio
 import logging
 import os
-import re
-import tempfile
-import uuid
-from typing import List, Tuple
+from typing import Any
 
-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")
-
-RENDER_SCALE = 2.0
-PAD_LEFT_FRAC = 0.035
-PAD_RIGHT_FRAC = 0.10
-PAD_TOP_FRAC = 0.018
-PAD_BOTTOM_FRAC = 0.018
-
-ENABLE_CONTRAST = True
-CONTRAST_FACTOR = 1.18
-ENABLE_UNSHARP = True
-UNSHARP_RADIUS = 0.78
-UNSHARP_PERCENT = 76
-UNSHARP_THRESHOLD = 1
-
-PAGE_PNG_COMPRESS_LEVEL = 3
-MAX_IMAGE_SIDE = 1568
-MAX_NEW_TOKENS = 3000
-
-# ── Model singleton ───────────────────────────────────────────────────────────
-_model = None
-_processor = None
-
-
-def _load_model():
-    global _model, _processor
-    if _model is not None:
-        return _model, _processor
-
-    import torch
-    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
-    )
-    _model = AutoModelForImageTextToText.from_pretrained(
-        MERGED_MODEL_DIR,
-        dtype=torch.bfloat16,
-        device_map="auto",
-        trust_remote_code=True,
-    )
-    _model.eval()
-    log.info("Model loaded.")
-    return _model, _processor
-
-
-def _enhance_raster_for_ocr(img):
-    from PIL import ImageEnhance, ImageFilter
-
-    if ENABLE_CONTRAST:
-        img = ImageEnhance.Contrast(img).enhance(CONTRAST_FACTOR)
-    if ENABLE_UNSHARP:
-        img = img.filter(
-            ImageFilter.UnsharpMask(
-                radius=UNSHARP_RADIUS,
-                percent=UNSHARP_PERCENT,
-                threshold=UNSHARP_THRESHOLD,
-            )
-        )
-    return img
-
-
-def _resize_for_inference(img):
-    """Resize image preserving aspect ratio so longest side <= MAX_IMAGE_SIDE."""
-    from PIL import Image
-
-    w, h = img.size
-    longest = max(w, h)
-    if longest <= MAX_IMAGE_SIDE:
-        return img
-    ratio = MAX_IMAGE_SIDE / longest
-    new_size = (int(w * ratio), int(h * ratio))
-    return img.resize(new_size, Image.LANCZOS)
-
-
-def _detect_headers(header_line: str) -> List[str]:
-    """Extract column headers from a header line by splitting on 2+ spaces."""
-    parts = re.split(r'\s{2,}', header_line.strip())
-    headers = [p.strip() for p in parts if p.strip()]
-    return headers if headers else ["POSTING DATE", "DESCRIPTION", "AMOUNT"]
-
-
-def _plaintext_rows_to_table(text: str) -> str:
-    """
-    Convert plain-text transaction rows (no pipe boundaries) into proper
-    markdown tables. Detects headers dynamically from the document.
-    Works for any column layout, not just Date/Description/Amount.
-    """
-    lines = text.split('\n')
-    result = []
-    i = 0
-    # Pattern to detect a header line (2+ words separated by 2+ spaces)
-    header_re = re.compile(
-        r'^[A-Z][A-Z\s]{3,}$',
-    )
-    # Pattern for a transaction row starting with MM/DD
-    txn_re = re.compile(r'^\d{2}/\d{2}\s+\S')
-    # Pattern to extract amount at end of line
-    amt_re = re.compile(r'^(.*?)\s+([\d,]+\.\d{2})$')
-
-    while i < len(lines):
-        line = lines[i]
-        stripped = line.strip()
-
-        # Already a proper markdown table row — pass through unchanged
-        if stripped.startswith('|'):
-            result.append(line)
-            i += 1
-            continue
-
-        # Detect start of a plain-text transaction block
-        if txn_re.match(stripped):
-            # Look back to find the most recent header line before this block
-            detected_headers = None
-            for j in range(len(result) - 1, max(len(result) - 5, -1), -1):
-                prev = result[j].strip()
-                if prev and not prev.startswith('|') and not prev.startswith('#'):
-                    candidate_headers = _detect_headers(prev)
-                    if len(candidate_headers) >= 2:
-                        detected_headers = candidate_headers
-                        # Remove the header line from result since we'll put it in table
-                        result.pop(j)
-                        break
-
-            if not detected_headers:
-                detected_headers = ["POSTING DATE", "DESCRIPTION", "AMOUNT"]
-
-            # Collect all transaction rows in this block
-            block = []
-            while i < len(lines):
-                l = lines[i].strip()
-                if not l:
-                    break
-                if l.startswith('|') or l.startswith('#'):
-                    break
-                if re.match(r'^Subtotal', l, re.IGNORECASE):
-                    break
-                if txn_re.match(l):
-                    block.append(l)
-                elif block:
-                    # Continuation line — merge into previous row
-                    block[-1] += ' ' + l
-                i += 1
-
-            if block:
-                # Build markdown table with detected headers
-                num_cols = len(detected_headers)
-                result.append('| ' + ' | '.join(detected_headers) + ' |')
-                result.append('| ' + ' | '.join(['---'] * num_cols) + ' |')
+try:
+    _orig_close = asyncio.BaseEventLoop.close
 
-                for row in block:
-                    if num_cols == 3:
-                        # Try to split into date / description / amount
-                        m = amt_re.match(row)
-                        if m:
-                            rest = m.group(1)
-                            amount = m.group(2)
-                            date_m = re.match(r'^(\d{2}/\d{2})\s+(.*)', rest)
-                            if date_m:
-                                result.append(
-                                    f'| {date_m.group(1)} | {date_m.group(2)} | {amount} |'
-                                )
-                                continue
-                        # No amount found — put full row in description, leave amount blank
-                        date_m = re.match(r'^(\d{2}/\d{2})\s+(.*)', row)
-                        if date_m:
-                            result.append(
-                                f'| {date_m.group(1)} | {date_m.group(2)} | |'
-                            )
-                        else:
-                            result.append(f'| | {row} | |')
-                    else:
-                        # For non-3-column tables just put full row as single cell
-                        result.append(f'| {row} |')
-            continue
-
-        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
-    from PIL import Image
-
-    model, processor = _load_model()
-
-    img = Image.open(image_path).convert("RGB")
-    img = _resize_for_inference(img)
-
-    fd, resized_path = tempfile.mkstemp(suffix=".png")
-    os.close(fd)
-    try:
-        img.save(resized_path, "PNG")
-
-        messages = [{
-            "role": "user",
-            "content": [
-                {"type": "image", "url": resized_path},
-                {"type": "text", "text": "Document Parsing:"},
-            ],
-        }]
-
-        inputs = processor.apply_chat_template(
-            messages,
-            tokenize=True,
-            add_generation_prompt=True,
-            return_dict=True,
-            return_tensors="pt",
-        ).to(model.device)
-        inputs.pop("token_type_ids", None)
-
-        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,
-            )
-
-        result = processor.decode(
-            ids[0][inputs["input_ids"].shape[1]:],
-            skip_special_tokens=True,
-        )
-
-        result = _plaintext_rows_to_table(result)
-        return result.strip()
-
-    finally:
+    def _safe_close(self):
         try:
-            os.unlink(resized_path)
-        except Exception:
+            _orig_close(self)
+        except (ValueError, OSError):
             pass
 
+    asyncio.BaseEventLoop.close = _safe_close
+except Exception:
+    pass
 
-def render_pdf_pages_to_images(pdf_path: str) -> Tuple[List[str], List[int]]:
-    import pymupdf as fitz
-    from PIL import Image
+log = logging.getLogger("glmocr_raw_app")
+logging.basicConfig(level=logging.INFO)
 
-    doc = fitz.open(pdf_path)
-    page_images: List[str] = []
-    page_heights: List[int] = []
+GLMOCR_API_KEY = "cee1d52dd91a4ab591b3f6e105f8ad89.LgbQTECuzX0zrito"
 
-    for i in range(len(doc)):
-        page = doc[i]
-        pix = page.get_pixmap(
-            matrix=fitz.Matrix(RENDER_SCALE, RENDER_SCALE), alpha=False
-        )
+_parser = None
 
-        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 * PAD_RIGHT_FRAC)
-        pad_t = int(h * PAD_TOP_FRAC)
-        pad_b = int(h * PAD_BOTTOM_FRAC)
+def get_parser():
+    global _parser
+    if _parser is None:
+        from glmocr import GlmOcr
 
-        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
+        kwargs = {"mode": "maas"}
+        if GLMOCR_API_KEY:
+            kwargs["api_key"] = GLMOCR_API_KEY
+        _parser = GlmOcr(**kwargs)
+    return _parser
 
-        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)
 
-    doc.close()
-    return page_images, page_heights
+def _extract_markdown(result: Any) -> str:
+    if result is None:
+        return ""
+    if isinstance(result, list):
+        parts = []
+        for item in result:
+            md = getattr(item, "markdown_result", "")
+            if md:
+                parts.append(str(md).strip())
+        return "\n\n---page-separator---\n\n".join(p for p in parts if p).strip()
+    md = getattr(result, "markdown_result", "")
+    return str(md).strip() if md else ""
 
 
-def run_ocr(uploaded_file):
+def run_ocr_raw(uploaded_file):
     if uploaded_file is None:
         return "Please upload a file."
-
-    page_images: 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 = []
-        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
-
+        parser = get_parser()
+        result = parser.parse(path)
+        markdown = _extract_markdown(result)
+        return markdown or "(No content)"
     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:
-            try:
-                if (
-                    isinstance(p, str)
-                    and p.endswith(".png")
-                    and "glmocr_page_" in os.path.basename(p)
-                ):
-                    os.unlink(p)
-            except Exception:
-                pass
+        log.exception("run_ocr_raw failed: %s", e)
+        return f"Error: {e}\n\n{traceback.format_exc()}"
 
 
 def _create_gradio_demo():
     import gradio as gr
 
-    with gr.Blocks(title="GLM-OCR Fine-tuned") as demo:
-        gr.Markdown("# GLM-OCR (Fine-tuned)")
-        gr.Markdown("Upload a PDF or image to extract structured markdown content.")
+    with gr.Blocks(title="GLM-OCR Raw (no fixes)") as demo:
+        gr.Markdown("# GLM-OCR Raw (No Fixes)")
         file_in = gr.File(
             label="Upload PDF or image",
             file_types=[".pdf", ".png", ".jpg", ".jpeg", ".tiff", ".bmp"],
         )
         run_btn = gr.Button("Run OCR", variant="primary")
-        out = gr.Textbox(lines=40, label="Output (markdown)")
-        run_btn.click(fn=run_ocr, inputs=file_in, outputs=out)
+        out = gr.Textbox(lines=40, label="Raw markdown_result")
+        run_btn.click(fn=run_ocr_raw, inputs=file_in, outputs=out)
     return demo
 
 
 if __name__ == "__main__":
-    _load_model()
-    _create_gradio_demo().launch()
+    _create_gradio_demo().launch()