import re from itertools import accumulate from typing import TYPE_CHECKING, Dict, List, Optional, Union from transformers.utils import logging from transformers.feature_extraction_utils import BatchFeature from transformers.image_utils import ImageInput, is_valid_image, load_image from transformers.processing_utils import Unpack from transformers.tokenization_utils_base import TextInput, BatchEncoding # Import the *parent* class from transformers.models.idefics3.processing_idefics3 import ( Idefics3Processor, Idefics3ProcessorKwargs, is_url, is_image_or_image_url, get_image_prompt_string, ) if TYPE_CHECKING: from transformers.tokenization_utils_base import PreTokenizedInput logger = logging.get_logger(__name__) class SmolLMMProcessor(Idefics3Processor): """ A subclass of Idefics3Processor that adds an `allow_mismatch` argument to skip the 1:1 match check between # tokens and #images. """ def __call__( self, images: Union[ImageInput, List[ImageInput], List[List[ImageInput]]] = None, text: Union[TextInput, "PreTokenizedInput", List[TextInput], List["PreTokenizedInput"]] = None, audio=None, videos=None, image_seq_len: Optional[int] = None, allow_mismatch: bool = False, # <--- NEW ARG **kwargs: Unpack[Idefics3ProcessorKwargs], ) -> BatchEncoding: """ Process input for Idefics3. If `allow_mismatch=True`, we skip the error when # tokens != #images. See `Idefics3Processor.__call__` docstring for details on the other params. """ if text is None and images is None: raise ValueError("You must provide either `text` or `images`.") # Merge default keyword args for text/images output_kwargs = self._merge_kwargs( Idefics3ProcessorKwargs, tokenizer_init_kwargs=self.tokenizer.init_kwargs, **kwargs, ) # If no override for image_seq_len is passed, use the default self.image_seq_len image_seq_len = image_seq_len if image_seq_len is not None else self.image_seq_len # Count how many tokens are in each text sample n_images_in_text = [] if text is not None: if isinstance(text, str): text = [text] n_images_in_text = [sample.count(self.image_token.content) for sample in text] inputs = BatchFeature() # --------------------------------------------------------------- # If images are provided, do all the logic that normally raises a mismatch error. # We'll skip or warn if allow_mismatch is True. # --------------------------------------------------------------- if images is not None: # Flatten or interpret images if is_image_or_image_url(images): images = [[images]] elif isinstance(images, list) and is_image_or_image_url(images[0]): # Original code raises error if mismatch. We'll skip if allow_mismatch. if not allow_mismatch and text is not None and sum(n_images_in_text) != len(images): raise ValueError( f"The total number of tokens in the prompts should match " f"the number of images. Found {sum(n_images_in_text)} tokens " f"but {len(images)} images." ) else: if text is not None and sum(n_images_in_text) != len(images): logger.warning( "Mismatch # tokens vs. #images, but allow_mismatch=True => continuing." ) # Re-group images to match text samples # if text is not None: # cumsum_images_in_text = [0] + list(accumulate(n_images_in_text)) # images = [ # images[cumsum_images_in_text[i] : cumsum_images_in_text[i + 1]] # for i in range(len(n_images_in_text)) # ] if text is not None: # Calculate frames per token total_images = len(images) total_tokens = sum(n_images_in_text) if total_images > total_tokens and total_images % total_tokens == 0: frames_per_token = total_images // total_tokens # Create new grouping that preserves consecutive frames new_images = [] for i in range(len(n_images_in_text)): start_idx = i * frames_per_token * n_images_in_text[i] end_idx = start_idx + (frames_per_token * n_images_in_text[i]) new_images.append(images[start_idx:end_idx]) images = new_images else: # Original regrouping logic for other cases cumsum_images_in_text = [0] + list(accumulate(n_images_in_text)) images = [ images[cumsum_images_in_text[i] : cumsum_images_in_text[i + 1]] for i in range(len(n_images_in_text)) ] else: images = [images] elif ( not isinstance(images, list) and not isinstance(images[0], list) and not is_image_or_image_url(images[0][0]) ): raise ValueError("Invalid input images. Provide image or list of images or list of list of images.") n_images_in_images = [len(sample) for sample in images] # Actually load images if they are URLs images = [[load_image(im) if is_url(im) else im for im in sample] for sample in images] # Let the parent's image_processor handle shape, resizing, etc. image_inputs = self.image_processor(images, **output_kwargs["images_kwargs"]) inputs.update(image_inputs) # If we have text, handle expansions if text is not None: if not allow_mismatch and n_images_in_images != n_images_in_text: raise ValueError( f"Mismatch in #images vs # tokens. We found {n_images_in_text} tokens " f"but have {n_images_in_images} images in each batch." ) else: if n_images_in_images != n_images_in_text: logger.warning( "Mismatch in #images vs # tokens, but allow_mismatch=True => continuing." ) # Rows/cols for expanded patch tokens image_rows = inputs.pop("rows", [[0] * len(text)]) image_cols = inputs.pop("cols", [[0] * len(text)]) fake_image_token = self.fake_image_token.content image_token = self.image_token.content global_img_token = self.global_image_tag prompt_strings = [] for sample, sample_rows, sample_cols in zip(text, image_rows, image_cols): image_prompt_strings = [] for n_rows, n_cols in zip(sample_rows, sample_cols): image_prompt_string = get_image_prompt_string( n_rows, n_cols, image_seq_len, image_token=image_token, fake_token_around_image=fake_image_token, global_img_token=global_img_token, ) image_prompt_strings.append(image_prompt_string) split_sample = sample.split(image_token) if len(split_sample) == 0: raise ValueError("Expected token in text, found none.") # Insert expansions for each placeholder combined = split_sample[0] # for i, image_prompt_string in enumerate(image_prompt_strings): # combined += image_prompt_string + split_sample[i + 1] for i, split_subsample in enumerate(split_sample[1:]): combined += image_prompt_strings[i-1] + split_subsample prompt_strings.append(combined) # Now tokenize the text with expansions text_inputs = self.tokenizer(text=prompt_strings, **output_kwargs["text_kwargs"]) inputs.update(text_inputs) # ------------------------------------------------------------------- # If we have text only (no images) # ------------------------------------------------------------------- elif text is not None: # no images => zero tokens if any(n_images_in_text): raise ValueError( f"Found {sum(n_images_in_text)} tokens in text, but no images were passed." ) text_inputs = self.tokenizer(text=text, **output_kwargs["text_kwargs"]) inputs.update(text_inputs) return inputs # batch_decode, decode, model_input_names remain the same as parent # If you want them identical, no need to override them.