mlx-community/Molmo2-8B-4bit

1"""2Processor class for Molmo2.3"""4from typing import Optional, Union5import dataclasses6 7import numpy as np8 9from transformers.image_utils import ImageInput10from transformers.video_utils import VideoInput11from transformers.processing_utils import (12    Unpack,13    ProcessingKwargs,14    ProcessorMixin,15)16from transformers.feature_extraction_utils import BatchFeature17from transformers.tokenization_utils_base import TextInput, PreTokenizedInput18from transformers.utils import logging19 20from transformers import AutoTokenizer21from .image_processing_molmo2 import Molmo2ImagesKwargs, Molmo2ImageProcessor22from .video_processing_molmo2 import Molmo2VideoProcessorKwargs, Molmo2VideoProcessor23 24 25logger = logging.get_logger(__name__)26 27 28# Special tokens, these should be present in any tokenizer we use since the preprocessor uses them29IMAGE_PATCH_TOKEN = f"<im_patch>"  # Where to insert high-res tokens30IMAGE_LOW_RES_TOKEN = f"<im_low>"  # Where to insert low-res tokens31IM_START_TOKEN = f"<im_start>"32LOW_RES_IMAGE_START_TOKEN = f"<low_res_im_start>"33FRAME_START_TOKEN = f"<frame_start>"34IM_END_TOKEN = f"<im_end>"35FRAME_END_TOKEN= f"<frame_end>"36IM_COL_TOKEN = f"<im_col>"37IMAGE_PROMPT = "<|image|>"38VIDEO_PROMPT = "<|video|>"39 40IMAGE_TOKENS = [41    IMAGE_PATCH_TOKEN,42    IM_COL_TOKEN,43    IM_START_TOKEN,44    LOW_RES_IMAGE_START_TOKEN,45    FRAME_START_TOKEN,46    IM_END_TOKEN,47    FRAME_END_TOKEN,48    IMAGE_LOW_RES_TOKEN,49]50 51 52class Molmo2ProcessorKwargs(ProcessingKwargs, total=False):53    """Molmo2 processor kwargs"""54    images_kwargs: Molmo2ImagesKwargs55    videos_kwargs: Molmo2VideoProcessorKwargs56    _defaults = {57        "text_kwargs": {58            "padding": False,59            "return_mm_token_type_ids": True,60        },61        "videos_kwargs": {"return_metadata": True},62    }63 64 65class Molmo2Processor(ProcessorMixin):66    attributes = ["image_processor", "video_processor", "tokenizer"]67    optional_attributes = [68        "chat_template",69        "time_mode",70        "image_use_col_tokens",71        "use_single_crop_col_tokens",72        "use_single_crop_start_token",73        "video_use_col_tokens",74        "use_frame_special_tokens",75    ]76    image_processor_class = "AutoImageProcessor"77    video_processor_class = "AutoVideoProcessor"78    tokenizer_class = "AutoTokenizer"79 80    def __init__(81        self,82        image_processor: Molmo2ImageProcessor = None,83        video_processor: Molmo2VideoProcessor = None,84        tokenizer: AutoTokenizer = None,85        chat_template: Optional[str] = None,86        image_use_col_tokens: Optional[bool] = True,87        use_single_crop_col_tokens: Optional[bool] = None,88        use_single_crop_start_token: Optional[bool] = True,89        video_use_col_tokens: Optional[bool] = False,90        use_frame_special_tokens: Optional[bool] = True,91        **kwargs92    ) -> None:93        super().__init__(94            image_processor,95            video_processor,96            tokenizer,97            chat_template=chat_template,98            image_use_col_tokens=image_use_col_tokens,99            use_single_crop_col_tokens=use_single_crop_col_tokens,100            use_single_crop_start_token=use_single_crop_start_token,101            video_use_col_tokens=video_use_col_tokens,102            use_frame_special_tokens=use_frame_special_tokens,103        )104 105        self.image_placeholder_token = IMAGE_PROMPT106        self.video_placeholder_token = VIDEO_PROMPT107        self.image_token_ids = [108            tokenizer.convert_tokens_to_ids(token)109            for token in IMAGE_TOKENS110        ]111 112    def get_image_tokens(self, image_grid: np.ndarray):113        resized_h, resized_w, height, width = image_grid114        per_row = np.full(width, IMAGE_PATCH_TOKEN)115        if self.image_use_col_tokens:116            per_row = np.concatenate([per_row, [IM_COL_TOKEN]], 0)117        joint = [118            [IM_START_TOKEN],119            np.tile(per_row, [height]),120            [IM_END_TOKEN],121        ]122        per_row = np.full(resized_w, IMAGE_PATCH_TOKEN)123        use_single_crop_col_tokens = (124            self.image_use_col_tokens125            if self.use_single_crop_col_tokens is None126            else self.use_single_crop_col_tokens127        )128        image_start_token = (129            LOW_RES_IMAGE_START_TOKEN130            if self.use_single_crop_start_token131            else IM_START_TOKEN132        )133        if use_single_crop_col_tokens:134            per_row = np.concatenate([per_row, [IM_COL_TOKEN]], 0)135        joint = [136            [image_start_token],137            np.tile(per_row, [resized_h]),138            [IM_END_TOKEN],139        ] + joint140 141        return np.concatenate(joint)142    143    def get_video_string(144        self,145        video_grid: np.ndarray,146        timestamps: np.ndarray,147    ):  148        if self.use_frame_special_tokens:149            start_token_id = FRAME_START_TOKEN150            end_token_id = FRAME_END_TOKEN151        else:152            start_token_id = IM_START_TOKEN153            end_token_id = IM_END_TOKEN154        155        num_frames, h, w = video_grid156        video_string: str = ""157        for frame_idx, frame_time in enumerate(timestamps):158            # `per-frame-compact` time mode159            prev_space = " " if frame_idx > 0 else ""160            frame_prefix = prev_space + f"{frame_time:.1f} " # explicit whitespace before/after image tokens161 162            video_string += frame_prefix163            per_row = np.full(w, IMAGE_PATCH_TOKEN)164            if self.video_use_col_tokens:165                per_row = np.concatenate([per_row, [IM_COL_TOKEN]], 0)166            extra_tokens = np.tile(per_row, [h])167            video_tokens = [168                [start_token_id],169                extra_tokens,170                [end_token_id],171            ]172            video_string += "".join(np.concatenate(video_tokens, 0))173 174        return video_string175 176    def insert_bos(177        self,178        input_ids: np.ndarray,179        attention_mask: np.ndarray,180        bos_token_id: int,181        pad_token_id: int,182    ):183        """184        Args:185            input_ids: [B, S] array with left padding186            attention_mask: [B, S] array (0 for pad, 1 for valid)187            bos_token_id: int188            pad_token_id: int189        Returns:190            input_ids_out: [B, S] or [B, S+1] array with bos inserted if needed191            attention_mask_out: same shape as input_ids_out192        """193 194        need_to_expand = len(input_ids.shape) == 1195        if need_to_expand:196            input_ids = input_ids[None, :]197            attention_mask = attention_mask[None, :]198 199        B, S = input_ids.shape200 201        # Handle zero-length sequence202        if S == 0:203            new_input_ids = np.full((B, 1), bos_token_id, dtype=input_ids.dtype)204            new_attention_mask = np.ones((B, 1), dtype=attention_mask.dtype)205            if need_to_expand:206                new_input_ids = new_input_ids[0]207                new_attention_mask = new_attention_mask[0]208            return new_input_ids, new_attention_mask209 210        first_valid_index = (attention_mask == 1).argmax(axis=-1)  # [B]211        bos_already_present = np.all(input_ids[np.arange(B), first_valid_index] == bos_token_id)212 213        if bos_already_present:214            if need_to_expand:215                input_ids = input_ids[0]216                attention_mask = attention_mask[0]217            return input_ids, attention_mask218        else:219            new_input_ids = np.full((B, S+1), pad_token_id, dtype=input_ids.dtype)220            new_attention_mask = np.zeros((B, S+1), dtype=attention_mask.dtype)221 222            src_idx = np.tile(np.arange(S), (B, 1))  # [B, S]223            valid_mask = src_idx >= first_valid_index[:, None]  # [B, S]224            tgt_idx = src_idx + 1  # shit right225            batch_idx = np.tile(np.arange(B)[:, None], (1, S))  # [B, S]226 227            # flatten valid_positions228            flat_vals = input_ids[valid_mask]229            flat_batch = batch_idx[valid_mask]230            flat_tgt = tgt_idx[valid_mask]231 232            new_input_ids[flat_batch, flat_tgt] = flat_vals233            new_attention_mask[flat_batch, flat_tgt] = 1234            235            insert_pos = first_valid_index236            new_input_ids[np.arange(B), insert_pos] = bos_token_id237            new_attention_mask[np.arange(B), insert_pos] = 1238 239            if need_to_expand:240                new_input_ids = new_input_ids[0]241                new_attention_mask = new_attention_mask[0]242 243            return new_input_ids, new_attention_mask244 245    def __call__(246        self,247        text: Union[TextInput, PreTokenizedInput, list[TextInput], list[PreTokenizedInput]] = None,248        images: ImageInput = None,249        videos: VideoInput = None,250        **kwargs: Unpack[Molmo2ProcessorKwargs],251    ) -> BatchFeature:252        """253 254        Args:255            text (`str`, `list[str]`, `list[list[str]]`):256                The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings257                (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set258                `is_split_into_words=True` (to lift the ambiguity with a batch of sequences).259            images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `list[PIL.Image.Image]`, `list[np.ndarray]`, `list[torch.Tensor]`):260                The image or batch of images to be prepared. Each image can be a PIL image, NumPy array or PyTorch261                tensor. Both channels-first and channels-last formats are supported.262            videos (`dict[str, Any]` or `list[dict[str, Any]]`):263                The video or batch of videos to be prepared. Each video can be a dictionary with the following keys:264                - `"frames"`: `np.ndarray` of shape (T, H, W, 3)265                - `"timestamps"`: `np.ndarray` of shape (T,)266                - `"sampled_fps"`: `float` (optional)267                - `"sampling_augmentation"`: `str` (optional)268            return_tensors (`str` or [`~utils.TensorType`], *optional*):269                If set, will return tensors of a particular framework. Acceptable values are:270                - `'tf'`: Return TensorFlow `tf.constant` objects.271                - `'pt'`: Return PyTorch `torch.Tensor` objects.272                - `'np'`: Return NumPy `np.ndarray` objects.273                - `'jax'`: Return JAX `jnp.ndarray` objects.274 275        Returns:276            `BatchFeature`: A [`BatchFeature`] with the following fields:277            - **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`.278            - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when279              `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not `None`).280            - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`.281            - **image_token_pooling** -- Indices of the patches in `image_grids` to pool for each token in `image_tokens`.282              Returned when `images` is not `None`.283            - **image_grids** -- Grids of images. Returned when `images` is not `None`.284            - **image_num_crops** -- Number of crops for each image. Returned when `images` is not `None`.285            - **pixel_values_videos** -- Pixel values of videos to be fed to a model. Returned when `videos` is not `None`.286            - **video_token_pooling** -- Indices of the patches in `video_grids` to pool for each token in `video_tokens`.287              Returned when `videos` is not `None`.288            - **video_grids** -- Grids of videos. Returned when `videos` is not `None`.289        """290 291        output_kwargs = self._merge_kwargs(292            Molmo2ProcessorKwargs,293            tokenizer_init_kwargs=self.tokenizer.init_kwargs,294            **kwargs,295        )296 297        if images is not None:298            image_inputs = self.image_processor(images, **output_kwargs["images_kwargs"])299            image_grids = image_inputs["image_grids"]300        else:301            image_inputs = {}302            image_grids = None303 304        if videos is not None:305            videos_inputs = self.video_processor(videos=videos, **output_kwargs["videos_kwargs"])306            video_grids = videos_inputs["video_grids"]307            # If user has not requested video metadata, pop it308            if "return_metadata" not in kwargs:309                video_metadata = videos_inputs.pop("video_metadata")310            else:311                video_metadata = videos_inputs["video_metadata"]312        else:313            videos_inputs = {}314            video_grids = None315 316        if not isinstance(text, list):317            text = [text]318        319        text = text.copy() # below lines change text in-place320 321        if image_grids is not None:322            index = 0323            for i in range(len(text)):324                num_images = text[i].count(self.image_placeholder_token)325                image_grids_i = image_grids[index:index+num_images]326                for image_grid in image_grids_i:327                    image_tokens = self.get_image_tokens(image_grid)328                    image_string = "".join(image_tokens)329                    text[i] = text[i].replace(self.image_placeholder_token, image_string, 1)330                index += num_images331        332        if video_grids is not None:333            index = 0334            for i in range(len(text)):335                num_videos = text[i].count(self.video_placeholder_token)336                assert num_videos in {0, 1}, "At most one video is supported for now"337                video_grids_i = video_grids[index:index+num_videos]338                metadata_i = video_metadata[index:index+num_videos]339                for video_grid, metadata in zip(video_grids_i, metadata_i):340                    video_string = self.get_video_string(341                        video_grid,342                        metadata.timestamps,343                    )344                    text[i] = text[i].replace(self.video_placeholder_token, video_string, 1)345                index += num_videos346 347        return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None)348        return_mm_token_type_ids = output_kwargs["text_kwargs"].pop("return_mm_token_type_ids", False)349        text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"])350 351        input_ids = text_inputs["input_ids"]352        attention_mask = text_inputs["attention_mask"]353 354        input_ids = np.array(input_ids)355        attention_mask = np.array(attention_mask)356        357        bos = self.tokenizer.bos_token_id or self.tokenizer.eos_token_id358        input_ids, attention_mask = self.insert_bos(359            input_ids, attention_mask, bos, self.tokenizer.pad_token_id360        )361 362        if return_mm_token_type_ids:363            image_tokens = np.array(self.image_token_ids).astype(input_ids.dtype)364            token_type_ids = np.any(input_ids[:, :, None] == image_tokens[None, None, :], axis=-1)365            text_inputs["token_type_ids"] = token_type_ids.tolist()366        367        text_inputs["input_ids"] = input_ids.tolist()368        text_inputs["attention_mask"] = attention_mask.tolist()369 370        return BatchFeature(371            data={**text_inputs, **image_inputs, **videos_inputs},372            tensor_type=return_tensors,373        )374 375    def post_process_image_text_to_text(376        self, generated_outputs, skip_special_tokens=True, clean_up_tokenization_spaces=False, **kwargs377    ):378        """379        Post-process the output of the model to decode the text.380 381        Args:382            generated_outputs (`torch.Tensor` or `np.ndarray`):383                The output of the model `generate` function. The output is expected to be a tensor of shape `(batch_size, sequence_length)`384                or `(sequence_length,)`.385            skip_special_tokens (`bool`, *optional*, defaults to `True`):386                Whether or not to remove special tokens in the output. Argument passed to the tokenizer's `batch_decode` method.387            clean_up_tokenization_spaces (`bool`, *optional*, defaults to `False`):388                Whether or not to clean up the tokenization spaces. Argument passed to the tokenizer's `batch_decode` method.389            **kwargs:390                Additional arguments to be passed to the tokenizer's `batch_decode method`.391 392        Returns:393            `list[str]`: The decoded text.394        """395        return self.tokenizer.batch_decode(396            generated_outputs,397            skip_special_tokens=skip_special_tokens,398            clean_up_tokenization_spaces=clean_up_tokenization_spaces,399            **kwargs,400        )401 402 403Molmo2Processor.register_for_auto_class()