File size: 4,648 Bytes
96f4123 d006ae4 70f54cc 96f4123 70f54cc 96f4123 9fc5c52 96f4123 4e3a057 96f4123 d006ae4 96f4123 d006ae4 70f54cc d006ae4 96f4123 bb9bc8a 96f4123 bb9bc8a 96f4123 bb9bc8a 96f4123 32bc132 96f4123 bb9bc8a 96f4123 | 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 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 | import os
import sys
import re
import spaces
import torch
from unittest.mock import patch
from transformers import AutoModel, AutoTokenizer
from transformers.dynamic_module_utils import get_imports as _original_get_imports
# Capture original BEFORE patching — importing inside the patched fn causes
# infinite recursion because it would resolve to the patched version.
def _strip_flash_attn(filename, *args, **kwargs):
return [i for i in _original_get_imports(filename, *args, **kwargs) if i != "flash_attn"]
# Block flash_attn at runtime — HF ZeroGPU has it installed, so without this
# the model's own "try: import flash_attn" succeeds and forces FA2 on.
sys.modules["flash_attn"] = None # None = ImportError on any "import flash_attn"
# ---------------------------------------------------------------------------
# Global Initialization — Nemotron Parse (Nvidia Spatial Expert)
# Specializes in dense printed documents: tables, grids, bounding boxes.
# Loaded into CPU memory on startup. GPU is claimed inside @spaces.GPU only.
# ---------------------------------------------------------------------------
print("Initializing Nemotron Parse (Nvidia)... This may take a moment.")
parser_model_id = "nvidia/NVIDIA-Nemotron-Parse-v1.2" # Verified correct ID
hf_token = os.environ.get("HF_TOKEN")
def _patch_flash_attn_imports(filename, *args, **kwargs):
from transformers.dynamic_module_utils import get_imports as _orig
return [imp for imp in _orig(filename, *args, **kwargs) if imp != "flash_attn"]
try:
with patch(
"transformers.dynamic_module_utils.get_imports",
side_effect=_strip_flash_attn
):
parser_tokenizer = AutoTokenizer.from_pretrained(
parser_model_id,
trust_remote_code=True,
token=hf_token
)
parser_model = AutoModel.from_pretrained(
parser_model_id,
trust_remote_code=True,
torch_dtype=torch.bfloat16,
low_cpu_mem_usage=True,
token=hf_token
)
print("Nemotron Parse loaded successfully.")
except Exception as e:
print(f"Warning: Nemotron Parse failed to load: {e}")
parser_tokenizer = None
parser_model = None
@spaces.GPU
def run_nemotron_parse(image_path: str) -> dict:
"""
Spatial Expert: Runs Nemotron Parse on a raw receipt image.
Best at: dense printed bills, hospital invoices, utility statements.
"""
if parser_model is None or parser_tokenizer is None:
return {"source": "nemotron_parse", "text": "", "confidence": 0.0}
try:
from PIL import Image as PILImage
from transformers import AutoProcessor
parser_model.to("cuda")
parser_model.eval()
image = PILImage.open(image_path).convert("RGB")
# Nemotron Parse v1.2 requires a specific 4-token prompt format:
# <predict_bbox><predict_classes><output_markdown><predict_text_in_pic>
# The last token controls whether text-in-image is extracted.
prompt = "<predict_no_bbox><predict_no_classes><output_markdown><predict_text_in_pic>"
processor = AutoProcessor.from_pretrained(
parser_model_id, trust_remote_code=True,
token=os.environ.get("HF_TOKEN")
)
inputs = processor(text=prompt, images=image, return_tensors="pt").to("cuda")
with torch.no_grad():
out_ids = parser_model.generate(**inputs, max_new_tokens=600)
input_length = inputs["input_ids"].shape[-1]
raw = processor.decode(out_ids[0][input_length:], skip_special_tokens=True)
confidence = _score_confidence(raw)
return {"source": "nemotron_parse", "text": raw.strip(), "confidence": confidence}
except Exception as e:
print(f"Nemotron Parse inference error: {e}")
return {"source": "nemotron_parse", "text": "", "confidence": 0.0}
def _score_confidence(text: str) -> float:
"""
Scores how structured an extracted text looks.
Higher score = more price-like tokens = better quality output.
This is the 'arbiter' that picks the winner in the parallel engine.
"""
if not text or len(text.strip()) < 20:
return 0.0
# Count how many price-like patterns appear (e.g. $4.99 / ₹799 / 12.00)
price_hits = len(re.findall(r"[\$₹€£¥]\s*\d+[\.,]\d{2}|\d+[\.,]\d{2}", text))
# Count structured line markers (dashes, colons, tabs — signs of a table)
structure_hits = len(re.findall(r"[-:|\t]", text))
# Normalize to a 0–1 range, capped at 1.0
score = min(1.0, (price_hits * 0.15) + (structure_hits * 0.02))
return round(score, 3)
|