Aivanov_scan_ogc/pipeline/ocr_qwen.py

"""Wrapper singleton pour Qwen2.5-VL-3B.

Qwen2.5-VL-3B surpasse GLM-OCR sur les fiches OGC : extrait `dp_libelle`,
`praticien_conseil` (manuscrit !), `codage_reco.dp` et les 4 GHM/GHS là où
GLM-OCR échouait systématiquement. Un poil plus rapide aussi (3s vs 4s/page).

Coût : ~7 Go VRAM en bfloat16 (GLM-OCR = 2.2 Go) → tient sur RTX 5070 12 Go.
"""
import time
from pathlib import Path
import torch
from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration
from qwen_vl_utils import process_vision_info

MODEL_PATH = "Qwen/Qwen2.5-VL-3B-Instruct"


class QwenVLOCR:
    """Charge Qwen2.5-VL-3B une fois, réutilise le modèle pour toutes les pages."""
    _instance = None

    def __new__(cls):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
            cls._instance._init_model()
        return cls._instance

    def _init_model(self):
        t0 = time.time()
        import os as _os

        # max_pixels limite le nombre de patches visuels pour éviter l'OOM
        # sur images 300 dpi (2481x3509). ~800 patches = équilibre qualité/VRAM,
        # tient confortablement dans ~5-6 Go même avec d'autres processus GPU
        # en arrière-plan. Configurable via env var QWEN_MAX_PIXELS (en patches).
        max_pixels = int(_os.environ.get("QWEN_MAX_PIXELS", 800)) * 28 * 28
        self.processor = AutoProcessor.from_pretrained(
            MODEL_PATH,
            min_pixels=256 * 28 * 28,
            max_pixels=max_pixels,
        )

        # Device : "auto" par défaut (GPU si dispo), "cpu" pour forcer le CPU
        # quand la VRAM est saturée par d'autres process. Configurable via
        # QWEN_DEVICE=cpu.
        device = _os.environ.get("QWEN_DEVICE", "auto").lower()
        if device == "cpu":
            # Sur CPU on cherche à maximiser le throughput :
            # 1. Utiliser tous les cores via torch.set_num_threads (set_num_threads
            #    prime sur OMP_NUM_THREADS pour les ops PyTorch natifs).
            # 2. Choisir bfloat16 si le CPU le supporte nativement (Zen 5,
            #    Zen 4, Intel Sapphire Rapids+ ont AVX-512 BF16). Sinon float32.
            n_threads = int(_os.environ.get("TORCH_NUM_THREADS", _os.cpu_count() or 8))
            torch.set_num_threads(n_threads)
            try:
                torch.set_num_interop_threads(n_threads)
            except RuntimeError:
                pass  # déjà initialisé, ignorer

            # Détection AVX-512 BF16 via /proc/cpuinfo (Linux)
            use_bf16 = False
            try:
                with open("/proc/cpuinfo") as f:
                    flags = f.read()
                use_bf16 = "avx512_bf16" in flags or "amx_bf16" in flags
            except Exception:
                pass
            dtype = torch.bfloat16 if use_bf16 else torch.float32

            self.model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
                MODEL_PATH,
                torch_dtype=dtype,
                device_map={"": "cpu"},
                low_cpu_mem_usage=True,
            )
            self.device_used = "cpu"
            self.cpu_threads = n_threads
            self.cpu_dtype = str(dtype).replace("torch.", "")
        else:
            self.model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
                MODEL_PATH,
                torch_dtype=torch.bfloat16,
                device_map="auto",
            )
            self.device_used = "cuda" if torch.cuda.is_available() else "cpu"
            self.cpu_threads = None
            self.cpu_dtype = None
        self.model.eval()
        self.load_time = time.time() - t0
        self.vram_gb = torch.cuda.memory_allocated() / 1e9 if torch.cuda.is_available() else 0.0

    def run(self, image_path: str | Path, prompt: str, max_new_tokens: int = 2048) -> dict:
        """Exécute Qwen-VL sur une image avec un prompt, retourne {text, elapsed_s}."""
        image_path = str(image_path)
        messages = [{
            "role": "user",
            "content": [
                {"type": "image", "image": image_path},
                {"type": "text", "text": prompt},
            ],
        }]
        t0 = time.time()
        text = self.processor.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True,
        )
        image_inputs, video_inputs = process_vision_info(messages)
        inputs = self.processor(
            text=[text], images=image_inputs, videos=video_inputs,
            padding=True, return_tensors="pt",
        ).to(self.model.device)

        with torch.no_grad():
            generated_ids = self.model.generate(**inputs, max_new_tokens=max_new_tokens)
        out_ids = generated_ids[:, inputs.input_ids.shape[1]:]
        output = self.processor.batch_decode(out_ids, skip_special_tokens=True)[0]
        # Libérer la VRAM allouée par l'inférence (utile quand d'autres
        # processus tournent en parallèle sur le même GPU)
        del inputs, generated_ids, out_ids
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        return {"text": output.strip(), "elapsed_s": time.time() - t0}