"""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": } 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 = ( "" ) _PARSE_BLOCK_RE = re.compile( r"(.*?)]+>", re.DOTALL, ) _PARSE_STRAY_TOKEN_RE = re.compile(r"]*>|]*>|]+>") 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()