redac/vision.py

"""Image -> text extraction with a local vision model (MiniCPM-V-4.5).

Runs in-Space on ZeroGPU via the @spaces.GPU decorator. The model reads a
document/ID image and returns the fields as plain text, which then flows
through the same PII detection + redaction core as typed text.

MiniCPM-V-4.5 (model_type "minicpmv") uses trust_remote_code + model.chat()
and runs on transformers 4.x. The weights are downloaded/loaded on CPU
*before* entering the GPU window so the 16GB download can't blow the
ZeroGPU time budget; the GPU call only moves the model to CUDA and generates.

Local dev: set REDAC_MOCK=1 to skip the model entirely and return a canned
extraction, so the Gradio UI runs on a laptop with no GPU.
"""

from __future__ import annotations

import os
from functools import lru_cache

MODEL_ID = "openbmb/MiniCPM-V-4_5"

EXTRACTION_PROMPT = (
    "You are a document data extractor. Read this image and transcribe ALL "
    "personal and sensitive information you can find. Output one field per "
    "line as 'field: value'. Include, when present: full name, date of birth, "
    "address, passport/ID/driver-license number, national ID or social "
    "security number, phone, email, and any account or card numbers. "
    "Transcribe values exactly as written. Do not invent fields."
)

_MOCK_EXTRACTION = (
    "full name: John A. Doe\n"
    "date of birth: 1985-04-12\n"
    "address: 221B Baker Street, London\n"
    "passport number: X1234567\n"
    "national id number: 123-45-6789\n"
    "email: john.doe@example.com\n"
    "phone: +49 151 23456789"
)


def _is_mock() -> bool:
    return os.environ.get("REDAC_MOCK", "").strip() in {"1", "true", "True"}


# ZeroGPU decorator; degrade to a no-op decorator when `spaces` is absent
# (local dev) so the module imports cleanly off-Space.
try:
    import spaces  # type: ignore

    _gpu = spaces.GPU(duration=120)
except Exception:  # pragma: no cover - local fallback
    def _gpu(fn):
        return fn


@lru_cache(maxsize=1)
def _load_model():
    """Download + load weights on CPU. Cached. Called outside the GPU window."""
    import torch
    from transformers import AutoModel, AutoTokenizer

    model = AutoModel.from_pretrained(
        MODEL_ID,
        trust_remote_code=True,
        attn_implementation="sdpa",
        torch_dtype=torch.bfloat16,
    ).eval()
    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
    return model, tokenizer


@_gpu
def _chat(image, prompt: str) -> str:
    model, tokenizer = _load_model()  # cached -> instant inside GPU window
    model = model.to("cuda")
    msgs = [{"role": "user", "content": [image, prompt]}]
    return model.chat(image=None, msgs=msgs, tokenizer=tokenizer, sampling=False)


def extract_text_from_image(image, prompt: str | None = None) -> str:
    """Return extracted field text from a PIL image. Honors REDAC_MOCK."""
    if image is None:
        return ""
    if _is_mock():
        return _MOCK_EXTRACTION
    _load_model()  # warm the cache on CPU before claiming the GPU
    return _chat(image.convert("RGB"), prompt or EXTRACTION_PROMPT).strip()