multixtract 0.1.1__py3-none-any.whl

multixtract/providers/qwen2vl.py

@@ -0,0 +1,158 @@
+"""Qwen2.5-VL vision provider (optional extra: ``pip install multixtract[qwen2vl]``).
+
+Runs Qwen2.5-VL fully offline via HuggingFace transformers — no cloud, no API
+key. It is the recommended local vision model for document understanding in 2025:
+it leads the 7B class on DocVQA, ChartQA, TextVQA, and OCR benchmarks, supports
+dynamic image resolutions natively, and has first-class transformers integration.
+
+A CUDA GPU with 16–24 GB VRAM is recommended for BF16. For CPU or low-VRAM
+environments use the 3B variant (``Qwen/Qwen2.5-VL-3B-Instruct``) or a 4-bit
+quantised GGUF checkpoint.
+
+Example::
+
+    from multixtract import extract_document
+    from multixtract.providers import Qwen2VLVisionModel
+
+    vision = Qwen2VLVisionModel()                     # default: 7B
+    _, images = extract_document("report.pdf")
+    for img in images:
+        r = vision.analyze(img["image_bytes"], img["ext"])
+        print(r.caption, "|", r.description)
+"""
+from __future__ import annotations
+
+import logging
+import threading
+from typing import Optional
+
+from ..interfaces import VisionModel, VisionResult
+from ..vision import DEFAULT_SYSTEM_PROMPT, parse_vision_response
+from ._utils import _infer_device, _open_image_rgb
+
+log = logging.getLogger("multixtract")
+
+
+class Qwen2VLVisionModel(VisionModel):
+    """Offline ``VisionModel`` backed by Qwen2.5-VL (transformers).
+
+    Recommended local model for document/chart/diagram understanding. Leads the
+    7B class on DocVQA, ChartQA, and TextVQA as of 2025. Drop-in replacement for
+    cloud providers — returns the same ``VisionResult`` structure.
+
+    Heavy dependencies (``torch`` / ``transformers`` / ``qwen-vl-utils``) are
+    imported lazily; simply importing this module never pulls them in.
+    """
+
+    def __init__(
+        self,
+        model_id: str = "Qwen/Qwen2.5-VL-7B-Instruct",
+        system_prompt: str = DEFAULT_SYSTEM_PROMPT,
+        max_new_tokens: int = 512,
+        device: Optional[str] = None,
+        torch_dtype: str = "bfloat16",
+        load_in_4bit: bool = False,
+        model=None,
+        processor=None,
+    ) -> None:
+        """Load a Qwen2.5-VL model.
+
+        Args:
+            model_id: HuggingFace model id. Recommended checkpoints:
+                ``Qwen/Qwen2.5-VL-7B-Instruct`` (default, best accuracy),
+                ``Qwen/Qwen2.5-VL-3B-Instruct`` (lower VRAM / faster CPU).
+            system_prompt: Instruction sent with every image. Defaults to the
+                shared CAPTION/OCR_TEXT/DESCRIPTION prompt.
+            max_new_tokens: Generation budget per image.
+            device: ``"cuda"`` or ``"cpu"``. Auto-detected when omitted.
+            torch_dtype: dtype for GPU (``"bfloat16"`` recommended for Qwen2.5-VL).
+            load_in_4bit: 4-bit quantised load (needs ``bitsandbytes``); halves
+                VRAM — useful on 8–12 GB cards.
+            model / processor: Pre-built instances to reuse (skips loading).
+        """
+        self.model_id = model_id
+        self.system_prompt = system_prompt
+        self.max_new_tokens = max_new_tokens
+        self._lock = threading.Lock()
+
+        if model is not None and processor is not None:
+            self._model = model
+            self._processor = processor
+            self.device = device or _infer_device()
+            return
+
+        try:
+            import torch
+            from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration
+        except ImportError as exc:
+            raise ImportError(
+                "Qwen2.5-VL support requires transformers + torch: "
+                "pip install 'multixtract[qwen2vl]'"
+            ) from exc
+
+        self.device = device or _infer_device()
+
+        load_kwargs: dict = {"low_cpu_mem_usage": True}
+        if load_in_4bit:
+            load_kwargs["load_in_4bit"] = True
+            load_kwargs["device_map"] = "auto"  # required by bitsandbytes 4-bit
+        else:
+            dtype = getattr(torch, torch_dtype, torch.bfloat16)
+            load_kwargs["torch_dtype"] = dtype
+            load_kwargs["device_map"] = self.device
+
+        self._processor = AutoProcessor.from_pretrained(model_id)
+        self._model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+            model_id, **load_kwargs
+        )
+        self._model.eval()
+
+    def analyze(
+        self,
+        image_bytes: bytes,
+        ext: str = "png",
+        width: int = 0,
+        height: int = 0,
+    ) -> VisionResult:
+        """Describe one image. Never raises — returns an empty result on failure."""
+        try:
+            import torch
+
+            with _open_image_rgb(image_bytes) as image:
+                messages = [
+                    {
+                        "role": "user",
+                        "content": [
+                            {"type": "image", "image": image},
+                            {"type": "text", "text": self.system_prompt},
+                        ],
+                    }
+                ]
+                with self._lock:
+                    text_prompt = self._processor.apply_chat_template(
+                        messages, tokenize=False, add_generation_prompt=True
+                    )
+                    inputs = self._processor(
+                        text=[text_prompt],
+                        images=[image],
+                        return_tensors="pt",
+                        padding=True,
+                    )
+                    inputs = {k: v.to(self._model.device) for k, v in inputs.items()}
+                    with torch.inference_mode():
+                        output_ids = self._model.generate(
+                            **inputs,
+                            max_new_tokens=self.max_new_tokens,
+                            do_sample=False,
+                        )
+                    # Strip the input tokens from the output
+                    input_len = inputs["input_ids"].shape[1]
+                    generated = output_ids[:, input_len:]
+                    text = self._processor.batch_decode(
+                        generated, skip_special_tokens=True, clean_up_tokenization_spaces=False
+                    )[0]
+
+            return parse_vision_response(text)
+        except Exception:  # noqa: BLE001
+            log.warning("Qwen2VL analyze failed", exc_info=True)
+            return VisionResult()

multixtract/providers/smolvlm.py

@@ -0,0 +1,152 @@
+"""SmolVLM vision provider (optional extra: ``pip install multixtract[smolvlm]``).
+
+Runs SmolVLM fully offline via HuggingFace transformers — no cloud, no API key.
+At 2.2B parameters it is the recommended CPU-friendly option, offering meaningfully
+better DocVQA and ChartQA accuracy while still running on CPU
+without impractical wait times.
+
+Uses the standard ``AutoProcessor`` + ``AutoModelForVision2Seq`` API — no
+``trust_remote_code`` required.
+
+Example::
+
+    from multixtract import extract_document
+    from multixtract.providers import SmolVLMVisionModel
+
+    vision = SmolVLMVisionModel()          # downloads ~4 GB on first use
+    _, images = extract_document("report.pdf")
+    for img in images:
+        r = vision.analyze(img["image_bytes"], img["ext"])
+        print(r.caption, "|", r.ocr_text)
+"""
+from __future__ import annotations
+
+import logging
+import threading
+from typing import Optional
+
+from ..interfaces import VisionModel, VisionResult
+from ..vision import DEFAULT_SYSTEM_PROMPT, parse_vision_response
+from ._utils import _infer_device, _open_image_rgb
+
+log = logging.getLogger("multixtract")
+
+
+class SmolVLMVisionModel(VisionModel):
+    """Offline ``VisionModel`` backed by SmolVLM-2.2B-Instruct.
+
+    Best choice for CPU-only or memory-constrained environments. At 2.2B
+    parameters it significantly outperforms competing models on document and chart
+    benchmarks while remaining practical on CPU. No ``trust_remote_code``
+    required — uses the standard transformers vision-language API.
+
+    Heavy dependencies (``torch`` / ``transformers``) are imported lazily;
+    simply importing this module never pulls them in.
+    """
+
+    def __init__(
+        self,
+        model_id: str = "HuggingFaceTB/SmolVLM-2.2B-Instruct",
+        system_prompt: str = DEFAULT_SYSTEM_PROMPT,
+        max_new_tokens: int = 512,
+        device: Optional[str] = None,
+        torch_dtype: str = "bfloat16",
+        load_in_4bit: bool = False,
+        model=None,
+        processor=None,
+    ) -> None:
+        """Load SmolVLM.
+
+        Args:
+            model_id: HuggingFace model id. Defaults to
+                ``HuggingFaceTB/SmolVLM-2.2B-Instruct``.
+                ``HuggingFaceTB/SmolVLM-500M-Instruct`` is available for
+                extremely constrained environments (lower accuracy).
+            system_prompt: Instruction sent with every image. Defaults to the
+                shared, parseable CAPTION/OCR_TEXT/DESCRIPTION prompt.
+            max_new_tokens: Generation budget per image.
+            device: ``"cuda"`` or ``"cpu"``. Auto-detected when omitted.
+            torch_dtype: dtype name used on GPU (``"bfloat16"`` recommended).
+            load_in_4bit: 4-bit quantised load (needs ``bitsandbytes``).
+            model / processor: Pre-built instances to reuse (skips loading).
+        """
+        self.model_id = model_id
+        self.system_prompt = system_prompt
+        self.max_new_tokens = max_new_tokens
+
+        self._lock = threading.Lock()
+
+        if model is not None and processor is not None:
+            self._model, self._processor = model, processor
+            self.device = device or _infer_device()
+            return
+
+        try:
+            import torch
+            from transformers import AutoModelForVision2Seq, AutoProcessor
+        except ImportError as exc:
+            raise ImportError(
+                "SmolVLM support requires transformers + torch: "
+                "pip install 'multixtract[smolvlm]'"
+            ) from exc
+
+        self.device = device or _infer_device()
+        load_kwargs = {"low_cpu_mem_usage": True}
+        if self.device == "cuda":
+            load_kwargs["torch_dtype"] = getattr(torch, torch_dtype, torch.bfloat16)
+        else:
+            load_kwargs["torch_dtype"] = torch.float32
+        if load_in_4bit:
+            load_kwargs["load_in_4bit"] = True
+            load_kwargs["device_map"] = "auto"  # required by bitsandbytes 4-bit
+        else:
+            load_kwargs["device_map"] = self.device
+
+        self._processor = AutoProcessor.from_pretrained(model_id)
+        self._model = AutoModelForVision2Seq.from_pretrained(model_id, **load_kwargs)
+        self._model.eval()
+
+    def analyze(
+        self,
+        image_bytes: bytes,
+        ext: str = "png",
+        width: int = 0,
+        height: int = 0,
+    ) -> VisionResult:
+        """Describe one image. Never raises — returns an empty result on failure."""
+        try:
+            import torch
+
+            with _open_image_rgb(image_bytes) as image:
+                messages = [
+                    {
+                        "role": "user",
+                        "content": [
+                            {"type": "image"},
+                            {"type": "text", "text": self.system_prompt},
+                        ],
+                    }
+                ]
+                with self._lock:
+                    input_text = self._processor.apply_chat_template(
+                        messages, add_generation_prompt=True
+                    )
+                    inputs = self._processor(
+                        images=[image], text=input_text, return_tensors="pt"
+                    ).to(self._model.device)
+                    with torch.inference_mode():
+                        output_ids = self._model.generate(
+                            **inputs,
+                            max_new_tokens=self.max_new_tokens,
+                            do_sample=False,
+                        )
+                    # Decode only the newly generated tokens (skip the prompt).
+                    prompt_len = inputs["input_ids"].shape[-1]
+                    text = self._processor.decode(
+                        output_ids[0][prompt_len:], skip_special_tokens=True
+                    )
+
+            return parse_vision_response(text)
+        except Exception:  # noqa: BLE001 — never break the caller/pipeline
+            log.warning("SmolVLM analyze failed", exc_info=True)
+            return VisionResult()

multixtract/providers/storage.py

@@ -0,0 +1,97 @@
+"""Storage providers.
+
+``LocalDiskStore`` is dependency-free (core). ``AzureBlobStore`` requires the
+``[azure]`` extra. Both satisfy the :class:`~multixtract.interfaces.BlobStore`
+protocol so the pipeline treats them identically.
+"""
+from __future__ import annotations
+
+import json
+import os
+from typing import Optional
+
+from ..interfaces import BlobStore
+
+# Separator-only JSON — ~30-40%% smaller than the default readable form.
+_COMPACT_SEPARATORS = (",", ":")
+
+
+def _dumps(obj, compact: bool) -> str:
+    if compact:
+        return json.dumps(obj, separators=_COMPACT_SEPARATORS, ensure_ascii=False)
+    return json.dumps(obj, indent=2, ensure_ascii=False)
+
+
+class LocalDiskStore(BlobStore):
+    """Persist outputs to the local filesystem under *root*."""
+
+    def __init__(self, root: str) -> None:
+        self.root = os.path.abspath(root)
+
+    def _full(self, path: str) -> str:
+        full = os.path.join(self.root, path)
+        os.makedirs(os.path.dirname(full), exist_ok=True)
+        return full
+
+    def put_bytes(self, path: str, data: bytes, content_type: str = "") -> str:
+        full = self._full(path)
+        with open(full, "wb") as fh:
+            fh.write(data)
+        return full
+
+    def put_json(self, path: str, obj: object, compact: bool = False) -> str:
+        full = self._full(path)
+        with open(full, "w", encoding="utf-8") as fh:
+            fh.write(_dumps(obj, compact))
+        return full
+
+    def exists(self, path: str) -> bool:
+        return os.path.exists(os.path.join(self.root, path))
+
+
+class AzureBlobStore(BlobStore):
+    """Persist outputs to an Azure Blob Storage container.
+
+    Provide either an account ``credential`` (e.g. a ClientSecretCredential or
+    DefaultAzureCredential) or a ready ``blob_service_client``. Never embed
+    secrets in code — pass credentials in from a secret store.
+    """
+
+    def __init__(
+        self,
+        container: str,
+        prefix: str = "",
+        account_url: Optional[str] = None,
+        credential=None,
+        blob_service_client=None,
+    ) -> None:
+        if blob_service_client is None:
+            from azure.storage.blob import BlobServiceClient
+            if not account_url:
+                raise ValueError("account_url is required when blob_service_client is not given")
+            blob_service_client = BlobServiceClient(account_url=account_url, credential=credential)
+        self._svc = blob_service_client
+        self.container = container
+        self.prefix = prefix.strip("/")
+
+    def _blob(self, path: str):
+        name = f"{self.prefix}/{path}" if self.prefix else path
+        return self._svc.get_blob_client(container=self.container, blob=name), name
+
+    def put_bytes(self, path: str, data: bytes, content_type: str = "") -> str:
+        from azure.storage.blob import ContentSettings
+        client, name = self._blob(path)
+        settings = ContentSettings(content_type=content_type) if content_type else None
+        client.upload_blob(data, overwrite=True, content_settings=settings)
+        return name
+
+    def put_json(self, path: str, obj: object, compact: bool = False) -> str:
+        return self.put_bytes(
+            path,
+            _dumps(obj, compact).encode("utf-8"),
+            content_type="application/json",
+        )
+
+    def exists(self, path: str) -> bool:
+        client, _ = self._blob(path)
+        return client.exists()

multixtract/py.typed

@@ -0,0 +1 @@
+

multixtract/vision.py

@@ -0,0 +1,120 @@
+"""Vendor-neutral vision helpers.
+
+These utilities are reused by concrete VisionModel providers:
+  * a default system prompt that yields CAPTION / OCR_TEXT / DESCRIPTION
+  * base64 data-URL encoding with a fast path + PIL resize fallback
+  * a parser that turns the structured response into a VisionResult
+
+No vendor SDK is imported here — only Pillow (already a core dependency).
+"""
+from __future__ import annotations
+
+import base64
+import io
+from typing import Dict
+
+from PIL import Image
+
+from .interfaces import VisionResult
+
+# Default prompt — providers may override. Produces a strict, parseable format.
+# Tuned for technical / engineering documents and works across vision backends
+# (GPT-4o, Azure OpenAI, local Llama 3.2 Vision). DESCRIPTION may span multiple lines —
+# parse_vision_response captures the full block.
+DEFAULT_SYSTEM_PROMPT = (
+    "You are a meticulous image-analysis assistant for technical and engineering "
+    "documents (reports, datasheets, test results, CAD/schematics, charts, and "
+    "tables). Analyze the single image and respond using EXACTLY the three "
+    "labelled sections below, in this order, with nothing before or after them:\n\n"
+    "CAPTION: <one concise sentence naming what the image is>\n"
+    "OCR_TEXT: <every piece of visible text verbatim — titles, labels, numbers "
+    "with units, axis ticks, legend entries, callouts, and table cells; separate "
+    "items with semicolons; write NONE if there is no text>\n"
+    "DESCRIPTION: <a thorough description: first state the image type (photograph "
+    "/ line chart / bar chart / scatter plot / schematic / circuit diagram / "
+    "engineering drawing / table / flowchart / screenshot), then the subject and "
+    "its components, the quantitative data or measurements shown, and any trends, "
+    "comparisons, anomalies, relationships, or engineering significance. Use "
+    "multiple sentences and do NOT invent details that are not visible.>"
+)
+
+# MIME types commonly supported natively by vision models via data URL.
+_SUPPORTED_MIME: Dict[str, str] = {
+    "png": "image/png", "jpeg": "image/jpeg", "jpg": "image/jpeg",
+    "gif": "image/gif", "webp": "image/webp",
+}
+
+_MAX_VISION_DIM = 2048  # Typical internal cap for "high" detail mode.
+
+
+def to_data_url(image_bytes: bytes, ext: str, width: int = 0, height: int = 0) -> str:
+    """Convert raw image bytes to a base64 data URL.
+
+    Fast path: when dimensions are known, the format is supported, and the
+    image is within the size cap, skip PIL entirely (pure base64 encode).
+    Slow path: open with PIL to resize (> max dim) or convert unsupported
+    formats to PNG.
+    """
+    mime = _SUPPORTED_MIME.get(ext.lower())
+
+    if mime and width > 0 and height > 0 and max(width, height) <= _MAX_VISION_DIM:
+        return f"data:{mime};base64,{base64.b64encode(image_bytes).decode()}"
+
+    img = Image.open(io.BytesIO(image_bytes))
+    try:
+        width, height = img.size
+        if max(width, height) > _MAX_VISION_DIM or mime is None:
+            orig = img
+            img = orig.convert("RGB")
+            orig.close()
+            if max(width, height) > _MAX_VISION_DIM:
+                img.thumbnail((_MAX_VISION_DIM, _MAX_VISION_DIM), Image.Resampling.LANCZOS)
+            png_buffer = io.BytesIO()
+            img.save(png_buffer, format="PNG")
+            image_bytes = png_buffer.getvalue()
+            mime = "image/png"
+    finally:
+        img.close()
+
+    return f"data:{mime};base64,{base64.b64encode(image_bytes).decode()}"
+
+
+def parse_vision_response(text: str) -> VisionResult:
+    """Parse a structured CAPTION/OCR_TEXT/DESCRIPTION response.
+
+    Each section may span multiple lines (common with local vision models);
+    a section runs until the next labelled marker. Falls back to using the whole
+    response as the description if the model did not follow the format.
+    """
+    result = VisionResult()
+    key_map = {"CAPTION": "caption", "OCR_TEXT": "ocr_text", "DESCRIPTION": "description"}
+    buffers: Dict[str, list] = {"caption": [], "ocr_text": [], "description": []}
+
+    current = None
+    seen: set = set()   # sections already opened — each is valid exactly once
+    for raw in (text or "").strip().splitlines():
+        line = raw.strip()
+        marker = None
+        for prefix, attr in key_map.items():
+            if attr not in seen and line.upper().startswith(f"{prefix}:"):
+                marker = attr
+                buffers[attr].append(line.split(":", 1)[1].strip())
+                break
+        if marker is not None:
+            current = marker
+            seen.add(marker)
+        elif current is not None and line:
+            buffers[current].append(line)
+
+    for attr, parts in buffers.items():
+        value = " ".join(p for p in parts if p).strip()
+        if value:
+            setattr(result, attr, value)
+
+    # Treat explicit "none" OCR sentinel as empty.
+    if result.ocr_text.strip().upper() in {"NONE", "N/A"}:
+        result.ocr_text = ""
+
+    if not result.description:
+        result.description = (text or "").strip()
+    return result