File size: 3,123 Bytes
988ede0 2414447 988ede0 2414447 988ede0 2414447 988ede0 2414447 988ede0 2414447 988ede0 2414447 988ede0 2414447 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 | """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()
|