extraction.py

"""Receipt extraction (PDF / image -> text) via OpenBMB MiniCPM-V-4.6.

Active OCR backend: MiniCPM-V-4.6 (native HF transformers, ~1.3B). Digital PDFs'
text layers are read directly (no OCR); only pages without a text layer, and
uploaded images, go through the vision model.

Vision runs on the Space/Modal GPU. The model loads at module scope (so ZeroGPU
forks share it) and is gated by LOAD_VISION, so text-only local dev needn't pull
the vision weights/deps (`LOAD_VISION=0`).

The earlier NVIDIA Nemotron-Parse implementation is preserved at the bottom
(unused) as an alternative backend.
"""

from __future__ import annotations

import os
import re

VISION_MODEL_ID = os.environ.get("VISION_MODEL_ID", "openbmb/MiniCPM-V-4.6")
LOAD_VISION = os.environ.get("LOAD_VISION", "1") == "1"
VISION_MAX_NEW_TOKENS = int(os.environ.get("VISION_MAX_NEW_TOKENS", "1024"))

OCR_PROMPT = (
    "Transcribe this receipt or financial statement verbatim as plain text. "
    "Include every line item with its price, plus the merchant, date, subtotal, "
    "tax, and total. Preserve the original order. Output only the transcription."
)

IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".webp", ".bmp", ".tif", ".tiff"}


def _device() -> str:
    import torch

    if torch.cuda.is_available():
        return "cuda"
    if torch.backends.mps.is_available():
        return "mps"
    return "cpu"


class VisionModel:
    """MiniCPM-V-4.6 OCR behind a single ``ocr(images) -> str`` call.

    Mirrors the official openbmb/MiniCPM-V-4.6 demo: native
    AutoModelForImageTextToText, SDPA attention, images passed as
    {"type": "image", "image": <PIL>} and the downsample_mode="16x" kwarg.
    """

    def __init__(self, model_id: str = VISION_MODEL_ID):
        import torch
        from transformers import AutoModelForImageTextToText, AutoProcessor

        self.device = _device()
        self.dtype = torch.float32 if self.device == "cpu" else torch.bfloat16
        print(f"[spend-elegy] loading vision model {model_id} on {self.device}")
        self.processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
        self.model = (
            AutoModelForImageTextToText.from_pretrained(
                model_id,
                torch_dtype=self.dtype,
                attn_implementation="sdpa",
                trust_remote_code=True,
            )
            .to(self.device)
            .eval()
        )

    def ocr(self, images) -> str:
        import torch

        texts = []
        for image in images:
            messages = [
                {
                    "role": "user",
                    "content": [
                        {"type": "image", "image": image},
                        {"type": "text", "text": OCR_PROMPT},
                    ],
                }
            ]
            inputs = self.processor.apply_chat_template(
                messages,
                add_generation_prompt=True,
                tokenize=True,
                return_dict=True,
                return_tensors="pt",
                enable_thinking=False,
                processor_kwargs={
                    "downsample_mode": "16x",
                    "max_slice_nums": 9,
                    "use_image_id": True,
                },
            ).to(self.model.device)
            for key, value in inputs.items():
                if isinstance(value, torch.Tensor) and torch.is_floating_point(value):
                    inputs[key] = value.to(self.dtype)
            outputs = self.model.generate(
                **inputs,
                max_new_tokens=VISION_MAX_NEW_TOKENS,
                do_sample=False,
                downsample_mode="16x",
            )
            new_tokens = outputs[0][inputs["input_ids"].shape[-1] :]
            texts.append(
                self.processor.tokenizer.decode(
                    new_tokens, skip_special_tokens=True
                ).strip()
            )
        return "\n\n".join(t for t in texts if t)


# Module-scope singleton (shared across ZeroGPU forks); gated for local text dev.
vision_model = VisionModel() if LOAD_VISION else None


def ocr_images(images) -> str:
    """OCR a list of PIL images via the vision model (runs on the GPU stage)."""
    if vision_model is None:
        raise RuntimeError(
            "Vision model not loaded (LOAD_VISION=0). Use the text/paste path, or "
            "unset LOAD_VISION to enable PDF/image OCR."
        )
    return vision_model.ocr(images)


def pdf_to_images(path: str):
    """Render each PDF page to a PIL RGB image (used for pages w/o a text layer)."""
    import fitz  # PyMuPDF
    from PIL import Image

    images = []
    with fitz.open(path) as doc:
        for page in doc:
            pix = page.get_pixmap(dpi=170)
            images.append(Image.frombytes("RGB", (pix.width, pix.height), pix.samples))
    return images


def extract_from_file(path: str):
    """Return ``(text, images_to_ocr)`` for a receipt file.

    - ``.txt``  : read text, no OCR.
    - ``.pdf``  : read each page's text layer; pages with none are rendered for OCR.
    - image     : queued for OCR.

    The OCR itself (``ocr_images``) runs in the GPU stage, not here.
    """
    ext = os.path.splitext(path)[1].lower()
    if ext == ".txt":
        with open(path, encoding="utf-8", errors="replace") as fh:
            return fh.read().strip(), []
    if ext == ".pdf":
        import fitz
        from PIL import Image

        text_parts, to_ocr = [], []
        with fitz.open(path) as doc:
            for page in doc:
                page_text = page.get_text().strip()
                if page_text:
                    text_parts.append(page_text)
                else:
                    pix = page.get_pixmap(dpi=170)
                    to_ocr.append(
                        Image.frombytes("RGB", (pix.width, pix.height), pix.samples)
                    )
        return "\n\n".join(text_parts), to_ocr
    if ext in IMAGE_EXTS:
        from PIL import Image

        return "", [Image.open(path).convert("RGB")]
    raise ValueError(f"Unsupported file type: {ext} (use .txt, .pdf, or an image)")


# ---------------------------------------------------------------------------
# Preserved earlier implementation: NVIDIA Nemotron-Parse OCR (unused).
# Kept for reference / as an alternative backend; not wired into the app. Its
# deps (albumentations, timm) are not in requirements.txt — add them to use it.
# ---------------------------------------------------------------------------
PARSE_TASK_PROMPT = (
    "</s><s><predict_bbox><predict_classes><output_markdown><predict_no_text_in_pic>"
)
_PARSE_BLOCK_RE = re.compile(
    r"<x_\d+(?:\.\d+)?><y_\d+(?:\.\d+)?>(.*?)<x_\d+(?:\.\d+)?><y_\d+(?:\.\d+)?><class_[^>]+>",
    re.DOTALL,
)
_PARSE_STRAY_TOKEN_RE = re.compile(r"<x_[^>]*>|<y_[^>]*>|<class_[^>]+>")


def parse_output_to_text(raw: str) -> str:
    """Turn raw Nemotron-Parse output into plain text (one block per line)."""
    blocks = [m.group(1).strip() for m in _PARSE_BLOCK_RE.finditer(raw)]
    blocks = [b for b in blocks if b]
    if blocks:
        return "\n".join(blocks)
    return _PARSE_STRAY_TOKEN_RE.sub(" ", raw).strip()