tools/vision.py

import os
import sys
import types
import spaces
import torch
from PIL import Image
from concurrent.futures import ThreadPoolExecutor, as_completed

# --- HACK TO BYPASS FLASH ATTENTION COMPILATION ---
if "flash_attn" not in sys.modules:
    import importlib.machinery
    mock_module = types.ModuleType("flash_attn")
    mock_module.__spec__ = importlib.machinery.ModuleSpec("flash_attn", None)
    mock_module.__version__ = "2.6.0"
    sys.modules["flash_attn"] = mock_module

# Use the CORRECT classes confirmed by official HuggingFace docs:
# AutoModelForImageTextToText + AutoProcessor (NOT AutoModelForCausalLM + AutoTokenizer)
from transformers import AutoModelForImageTextToText, AutoProcessor

# ---------------------------------------------------------------------------
# Global Initialization — MiniCPM-V 4.6 (OpenBMB Handwriting Expert)
# ---------------------------------------------------------------------------
print("Initializing MiniCPM-V 4.6 (OpenBMB Handwriting Expert)...")

model_id = "openbmb/MiniCPM-V-4.6"
hf_token = os.environ.get("HF_TOKEN")

try:
    vision_processor = AutoProcessor.from_pretrained(
        model_id,
        trust_remote_code=True,
        token=hf_token
    )
    vision_model = AutoModelForImageTextToText.from_pretrained(
        model_id,
        trust_remote_code=True,
        torch_dtype=torch.bfloat16,
        low_cpu_mem_usage=True,
        token=hf_token
    )
    print("MiniCPM-V 4.6 loaded successfully.")
except Exception as e:
    print(f"Warning: MiniCPM-V failed to load: {e}")
    vision_processor = None
    vision_model = None


@spaces.GPU
def _run_minicpm_vision(image_path: str) -> dict:
    """
    Handwriting Expert: Runs MiniCPM-V 4.6 on a raw receipt image.
    Uses the correct AutoModelForImageTextToText + AutoProcessor API.
    """
    if vision_model is None or vision_processor is None:
        return {"source": "minicpm_v", "text": "", "confidence": 0.0}

    try:
        vision_model.to("cuda")
        vision_model.eval()

        image = Image.open(image_path).convert("RGB")

        prompt = (
            "You are a financial data extractor. "
            "Read this receipt carefully. "
            "Extract every line item, its price, taxes, fees, discounts, and the total. "
            "Preserve the original order. Format as a clean structured list."
        )

        # The .chat() method expects a simple string content since image is passed separately
        messages = [
            {
                "role": "user",
                "content": prompt,
            }
        ]

        # Use the official .chat() method as recommended by MiniCPM-V docs
        res = vision_model.chat(
            image=image,
            msgs=messages,
            tokenizer=vision_processor.tokenizer if hasattr(vision_processor, "tokenizer") else vision_processor,
            sampling=False, # Use greedy decoding for receipt extraction
            max_new_tokens=600
        )

        # Some versions of MiniCPM-V return a tuple (res, context, _), others return a string
        if isinstance(res, tuple):
            res = res[0]

        confidence = _score_confidence(res)
        return {"source": "minicpm_v", "text": res.strip(), "confidence": confidence}

    except Exception as e:
        print(f"MiniCPM-V inference error: {e}")
        return {"source": "minicpm_v", "text": "", "confidence": 0.0}


def _score_confidence(text: str) -> float:
    """
    Scores how structured extracted text looks.
    Counts price-like patterns and structural markers.
    Used by the arbiter to pick the winner in the parallel engine.
    """
    import re
    if not text or len(text.strip()) < 20:
        return 0.0
    price_hits = len(re.findall(r"[\$₹€£¥]\s*\d+[\.,]\d{2}|\d+[\.,]\d{2}", text))
    structure_hits = len(re.findall(r"[-:|\t]", text))
    score = min(1.0, (price_hits * 0.15) + (structure_hits * 0.02))
    return round(score, 3)


def process_receipt_image(image_path: str) -> str:
    """
    The Parallel Vision Engine.
    Dispatches the raw image to both models simultaneously.
    Confidence scorer picks the winner.
    """
    from tools.parser import run_nemotron_parse

    results = {}

    with ThreadPoolExecutor(max_workers=2) as executor:
        futures = {
            executor.submit(_run_minicpm_vision, image_path): "minicpm_v",
            executor.submit(run_nemotron_parse, image_path): "nemotron_parse",
        }

        for future in as_completed(futures):
            try:
                result = future.result()
                results[result["source"]] = result
            except Exception as e:
                print(f"Vision engine worker error: {e}")

    if not results:
        return "Error: Both vision models failed to process the image."

    best = sorted(results.values(), key=lambda r: r["confidence"], reverse=True)[0]

    print(
        f"[Parallel Vision Engine] Winner: {best['source']} "
        f"(confidence={best['confidence']})"
    )

    if best["confidence"] == 0.0:
        fallbacks = [r["text"] for r in results.values() if r["text"]]
        if fallbacks:
            return fallbacks[0]
        return "Error: Could not extract text from the image. Please try a clearer photo."

    return best["text"]