import logging import sys import time from collections.abc import Iterable from pathlib import Path from typing import Any, Optional, Union import numpy as np from PIL.Image import Image from transformers import AutoModelForImageTextToText, AutoProcessor, GenerationConfig from docling.datamodel.accelerator_options import ( AcceleratorOptions, ) from docling.datamodel.base_models import VlmPrediction, VlmStopReason from docling.datamodel.pipeline_options_vlm_model import InlineVlmOptions from docling.models.base_model import BaseVlmModel from docling.models.utils.hf_model_download import ( HuggingFaceModelDownloadMixin, ) from docling.utils.accelerator_utils import decide_device _log = logging.getLogger(__name__) # Source code from https://huggingface.co/numind/NuExtract-2.0-8B def process_all_vision_info(messages, examples=None): """ Process vision information from both messages and in-context examples, supporting batch processing. Args: messages: List of message dictionaries (single input) OR list of message lists (batch input) examples: Optional list of example dictionaries (single input) OR list of example lists (batch) Returns: A flat list of all images in the correct order: - For single input: example images followed by message images - For batch input: interleaved as (item1 examples, item1 input, item2 examples, item2 input, etc.) - Returns None if no images were found """ try: from qwen_vl_utils import fetch_image, process_vision_info except ImportError: raise ImportError( "qwen-vl-utils is required for NuExtractTransformersModel. " "Please install it with: pip install qwen-vl-utils" ) from qwen_vl_utils import fetch_image, process_vision_info # Helper function to extract images from examples def extract_example_images(example_item): if not example_item: return [] # Handle both list of examples and single example examples_to_process = ( example_item if isinstance(example_item, list) else [example_item] ) images = [] for example in examples_to_process: if ( isinstance(example.get("input"), dict) and example["input"].get("type") == "image" ): images.append(fetch_image(example["input"])) return images # Normalize inputs to always be batched format is_batch = messages and isinstance(messages[0], list) messages_batch = messages if is_batch else [messages] is_batch_examples = ( examples and isinstance(examples, list) and (isinstance(examples[0], list) or examples[0] is None) ) examples_batch = ( examples if is_batch_examples else ([examples] if examples is not None else None) ) # Ensure examples batch matches messages batch if provided if examples and len(examples_batch) != len(messages_batch): if not is_batch and len(examples_batch) == 1: # Single example set for a single input is fine pass else: raise ValueError("Examples batch length must match messages batch length") # Process all inputs, maintaining correct order all_images = [] for i, message_group in enumerate(messages_batch): # Get example images for this input if examples and i < len(examples_batch): input_example_images = extract_example_images(examples_batch[i]) all_images.extend(input_example_images) # Get message images for this input input_message_images = process_vision_info(message_group)[0] or [] all_images.extend(input_message_images) return all_images if all_images else None class NuExtractTransformersModel(BaseVlmModel, HuggingFaceModelDownloadMixin): def __init__( self, enabled: bool, artifacts_path: Optional[Path], accelerator_options: AcceleratorOptions, vlm_options: InlineVlmOptions, ): self.enabled = enabled self.vlm_options = vlm_options if self.enabled: import torch self.device = decide_device( accelerator_options.device, supported_devices=vlm_options.supported_devices, ) _log.debug(f"Available device for NuExtract VLM: {self.device}") self.max_new_tokens = vlm_options.max_new_tokens self.temperature = vlm_options.temperature repo_cache_folder = vlm_options.repo_id.replace("/", "--") if artifacts_path is None: artifacts_path = self.download_models( repo_id=self.vlm_options.repo_id, revision=self.vlm_options.revision, ) elif (artifacts_path / repo_cache_folder).exists(): artifacts_path = artifacts_path / repo_cache_folder self.processor = AutoProcessor.from_pretrained( artifacts_path, trust_remote_code=vlm_options.trust_remote_code, use_fast=True, ) self.processor.tokenizer.padding_side = "left" self.vlm_model = AutoModelForImageTextToText.from_pretrained( artifacts_path, device_map=self.device, dtype=self.vlm_options.torch_dtype, _attn_implementation=( "flash_attention_2" if self.device.startswith("cuda") and accelerator_options.cuda_use_flash_attention2 else "sdpa" ), trust_remote_code=vlm_options.trust_remote_code, ) if sys.version_info < (3, 14): self.vlm_model = torch.compile(self.vlm_model) # type: ignore else: self.vlm_model.eval() # Load generation config self.generation_config = GenerationConfig.from_pretrained(artifacts_path) def process_images( self, image_batch: Iterable[Union[Image, np.ndarray]], prompt: Union[str, list[str]], ) -> Iterable[VlmPrediction]: """ Batched inference for NuExtract VLM using the specialized input format. Args: image_batch: Iterable of PIL Images or numpy arrays prompt: Either: - str: Single template used for all images - list[str]: List of templates (one per image, must match image count) """ import torch from PIL import Image as PILImage # Normalize images to RGB PIL pil_images: list[Image] = [] for img in image_batch: if isinstance(img, np.ndarray): if img.ndim == 3 and img.shape[2] in (3, 4): pil_img = PILImage.fromarray(img.astype(np.uint8)) elif img.ndim == 2: pil_img = PILImage.fromarray(img.astype(np.uint8), mode="L") else: raise ValueError(f"Unsupported numpy array shape: {img.shape}") else: pil_img = img if pil_img.mode != "RGB": pil_img = pil_img.convert("RGB") pil_images.append(pil_img) if not pil_images: return # Normalize templates (1 per image) if isinstance(prompt, str): templates = [prompt] * len(pil_images) else: if len(prompt) != len(pil_images): raise ValueError( f"Number of templates ({len(prompt)}) must match number of images ({len(pil_images)})" ) templates = prompt # Construct NuExtract input format inputs = [] for pil_img, template in zip(pil_images, templates): input_item = { "document": {"type": "image", "image": pil_img}, "template": template, } inputs.append(input_item) # Create messages structure for batch processing messages = [ [ { "role": "user", "content": [x["document"]], } ] for x in inputs ] # Apply chat template to each example individually texts = [ self.processor.tokenizer.apply_chat_template( messages[i], template=x["template"], tokenize=False, add_generation_prompt=True, ) for i, x in enumerate(inputs) ] # Process vision inputs using qwen-vl-utils image_inputs = process_all_vision_info(messages) # Process with the processor processor_inputs = self.processor( text=texts, images=image_inputs, padding=True, return_tensors="pt", **self.vlm_options.extra_processor_kwargs, ) processor_inputs = {k: v.to(self.device) for k, v in processor_inputs.items()} # Generate gen_kwargs = { **processor_inputs, "max_new_tokens": self.max_new_tokens, "generation_config": self.generation_config, **self.vlm_options.extra_generation_config, } if self.temperature > 0: gen_kwargs["do_sample"] = True gen_kwargs["temperature"] = self.temperature else: gen_kwargs["do_sample"] = False start_time = time.time() with torch.inference_mode(): from typing import cast generated_ids = cast(Any, self.vlm_model).generate(**gen_kwargs) generation_time = time.time() - start_time # Trim generated sequences input_len = processor_inputs["input_ids"].shape[1] trimmed_sequences = generated_ids[:, input_len:] # Decode with the processor/tokenizer decoded_texts: list[str] = self.processor.batch_decode( trimmed_sequences, skip_special_tokens=True, clean_up_tokenization_spaces=False, ) # Optional logging num_tokens = None if generated_ids.shape[0] > 0: # type: ignore # Todo: confirm num tokens is actually from first item, code was already like this num_tokens = int(generated_ids[0].shape[0]) _log.debug( f"Generated {num_tokens} tokens in {generation_time:.2f}s " f"for batch size {generated_ids.shape[0]}." # type: ignore ) for text in decoded_texts: # Apply decode_response to the output text decoded_text = self.vlm_options.decode_response(text) yield VlmPrediction( text=decoded_text, generation_time=generation_time, num_tokens=num_tokens, stop_reason=VlmStopReason.UNSPECIFIED, )