mlx-community/MolmoPoint-8B-fp16

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, AllKwargsForChatTemplate,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        use_low_res_token_for_global_crops: bool = False,92        **kwargs93    ) -> None:94        super().__init__(95            image_processor,96            video_processor,97            tokenizer,98            chat_template=chat_template,99            image_use_col_tokens=image_use_col_tokens,100            use_single_crop_col_tokens=use_single_crop_col_tokens,101            use_single_crop_start_token=use_single_crop_start_token,102            video_use_col_tokens=video_use_col_tokens,103            use_frame_special_tokens=use_frame_special_tokens,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        self.use_low_res_token_for_global_crops = use_low_res_token_for_global_crops112        self._patch_metadata = None113 114    def get_image_tokens(self, image_grid: np.ndarray):115        resized_h, resized_w, height, width = image_grid116        per_row = np.full(width, IMAGE_PATCH_TOKEN)117        if self.image_use_col_tokens:118            per_row = np.concatenate([per_row, [IM_COL_TOKEN]], 0)119        joint = [120            [IM_START_TOKEN],121            np.tile(per_row, [height]),122            [IM_END_TOKEN],123        ]124        if self.use_low_res_token_for_global_crops:125            per_row = np.full(resized_w, IMAGE_LOW_RES_TOKEN)126        else:127            per_row = np.full(resized_w, IMAGE_PATCH_TOKEN)128        use_single_crop_col_tokens = (129            self.image_use_col_tokens130            if self.use_single_crop_col_tokens is None131            else self.use_single_crop_col_tokens132        )133        image_start_token = (134            LOW_RES_IMAGE_START_TOKEN135            if self.use_single_crop_start_token136            else IM_START_TOKEN137        )138        if use_single_crop_col_tokens:139            per_row = np.concatenate([per_row, [IM_COL_TOKEN]], 0)140        joint = [141            [image_start_token],142            np.tile(per_row, [resized_h]),143            [IM_END_TOKEN],144        ] + joint145 146        return np.concatenate(joint)147    148    def get_video_string(149        self,150        video_grid: np.ndarray,151        timestamps: np.ndarray,152    ):  153        if self.use_frame_special_tokens:154            start_token_id = FRAME_START_TOKEN155            end_token_id = FRAME_END_TOKEN156        else:157            start_token_id = IM_START_TOKEN158            end_token_id = IM_END_TOKEN159        160        num_frames, h, w = video_grid161        video_string: str = ""162        for frame_idx, frame_time in enumerate(timestamps):163            # `per-frame-compact` time mode164            prev_space = " " if frame_idx > 0 else ""165            frame_prefix = prev_space + f"{frame_time:.1f} " # explicit whitespace before/after image tokens166 167            video_string += frame_prefix168            per_row = np.full(w, IMAGE_PATCH_TOKEN)169            if self.video_use_col_tokens:170                per_row = np.concatenate([per_row, [IM_COL_TOKEN]], 0)171            extra_tokens = np.tile(per_row, [h])172            video_tokens = [173                [start_token_id],174                extra_tokens,175                [end_token_id],176            ]177            video_string += "".join(np.concatenate(video_tokens, 0))178 179        return video_string180 181    def insert_bos(182        self,183        input_ids: np.ndarray,184        attention_mask: np.ndarray,185        bos_token_id: int,186        pad_token_id: int,187    ):188        """189        Args:190            input_ids: [B, S] array with left padding191            attention_mask: [B, S] array (0 for pad, 1 for valid)192            bos_token_id: int193            pad_token_id: int194        Returns:195            input_ids_out: [B, S] or [B, S+1] array with bos inserted if needed196            attention_mask_out: same shape as input_ids_out197        """198 199        need_to_expand = len(input_ids.shape) == 1200        if need_to_expand:201            input_ids = input_ids[None, :]202            attention_mask = attention_mask[None, :]203 204        B, S = input_ids.shape205 206        # Handle zero-length sequence207        if S == 0:208            new_input_ids = np.full((B, 1), bos_token_id, dtype=input_ids.dtype)209            new_attention_mask = np.ones((B, 1), dtype=attention_mask.dtype)210            if need_to_expand:211                new_input_ids = new_input_ids[0]212                new_attention_mask = new_attention_mask[0]213            return new_input_ids, new_attention_mask214 215        first_valid_index = (attention_mask == 1).argmax(axis=-1)  # [B]216        bos_already_present = np.all(input_ids[np.arange(B), first_valid_index] == bos_token_id)217 218        if bos_already_present:219            if need_to_expand:220                input_ids = input_ids[0]221                attention_mask = attention_mask[0]222            return input_ids, attention_mask223        else:224            new_input_ids = np.full((B, S+1), pad_token_id, dtype=input_ids.dtype)225            new_attention_mask = np.zeros((B, S+1), dtype=attention_mask.dtype)226 227            src_idx = np.tile(np.arange(S), (B, 1))  # [B, S]228            valid_mask = src_idx >= first_valid_index[:, None]  # [B, S]229            tgt_idx = src_idx + 1  # shit right230            batch_idx = np.tile(np.arange(B)[:, None], (1, S))  # [B, S]231 232            # flatten valid_positions233            flat_vals = input_ids[valid_mask]234            flat_batch = batch_idx[valid_mask]235            flat_tgt = tgt_idx[valid_mask]236 237            new_input_ids[flat_batch, flat_tgt] = flat_vals238            new_attention_mask[flat_batch, flat_tgt] = 1239            240            insert_pos = first_valid_index241            new_input_ids[np.arange(B), insert_pos] = bos_token_id242            new_attention_mask[np.arange(B), insert_pos] = 1243 244            if need_to_expand:245                new_input_ids = new_input_ids[0]246                new_attention_mask = new_attention_mask[0]247 248            return new_input_ids, new_attention_mask249 250    def __call__(251        self,252        text: Union[TextInput, PreTokenizedInput, list[TextInput], list[PreTokenizedInput]] = None,253        images: ImageInput = None,254        videos: VideoInput = None,255        return_pointing_metadata: bool = False,256        use_low_res_token_for_global_crops: bool = False,257        **kwargs: Unpack[Molmo2ProcessorKwargs],258    ) -> BatchFeature:259        """260 261        Args:262            text (`str`, `list[str]`, `list[list[str]]`):263                The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings264                (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set265                `is_split_into_words=True` (to lift the ambiguity with a batch of sequences).266            images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `list[PIL.Image.Image]`, `list[np.ndarray]`, `list[torch.Tensor]`):267                The image or batch of images to be prepared. Each image can be a PIL image, NumPy array or PyTorch268                tensor. Both channels-first and channels-last formats are supported.269            videos (`dict[str, Any]` or `list[dict[str, Any]]`):270                The video or batch of videos to be prepared. Each video can be a dictionary with the following keys:271                - `"frames"`: `np.ndarray` of shape (T, H, W, 3)272                - `"timestamps"`: `np.ndarray` of shape (T,)273                - `"sampled_fps"`: `float` (optional)274                - `"sampling_augmentation"`: `str` (optional)275            return_tensors (`str` or [`~utils.TensorType`], *optional*):276                If set, will return tensors of a particular framework. Acceptable values are:277                - `'tf'`: Return TensorFlow `tf.constant` objects.278                - `'pt'`: Return PyTorch `torch.Tensor` objects.279                - `'np'`: Return NumPy `np.ndarray` objects.280                - `'jax'`: Return JAX `jnp.ndarray` objects.281 282        Returns:283            `BatchFeature`: A [`BatchFeature`] with the following fields:284            - **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`.285            - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when286              `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not `None`).287            - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`.288            - **image_token_pooling** -- Indices of the patches in `image_grids` to pool for each token in `image_tokens`.289              Returned when `images` is not `None`.290            - **image_grids** -- Grids of images. Returned when `images` is not `None`.291            - **image_num_crops** -- Number of crops for each image. Returned when `images` is not `None`.292            - **pixel_values_videos** -- Pixel values of videos to be fed to a model. Returned when `videos` is not `None`.293            - **video_token_pooling** -- Indices of the patches in `video_grids` to pool for each token in `video_tokens`.294              Returned when `videos` is not `None`.295            - **video_grids** -- Grids of videos. Returned when `videos` is not `None`.296        """297        output_kwargs = self._merge_kwargs(298            Molmo2ProcessorKwargs,299            tokenizer_init_kwargs=self.tokenizer.init_kwargs,300            **kwargs,301        )302        patch_metadata = {}303        if images is not None:304            image_inputs = self.image_processor(images, **output_kwargs["images_kwargs"],305                                                return_pointing_metadata=return_pointing_metadata)306            if return_pointing_metadata:307                patch_metadata["token_pooling"] = image_inputs.pop("image_token_pooling_np")308                patch_metadata["subpatch_mapping"] = image_inputs.pop("subpatch_mapping")309                patch_metadata["image_sizes"] = image_inputs.pop("image_sizes")310            image_grids = image_inputs["image_grids"]311        else:312            image_inputs = {}313            image_grids = None314 315        if videos is not None:316            videos_inputs = self.video_processor(317                videos=videos, **output_kwargs["videos_kwargs"],318                return_pointing_metadata=return_pointing_metadata319            )320            if return_pointing_metadata:321                assert len(videos_inputs['video_metadata']) == 1322                vd_metadata = videos_inputs['video_metadata'][0]323                patch_metadata["token_pooling"] = videos_inputs.pop("video_token_pooling_np")324                patch_metadata["subpatch_mapping"] = videos_inputs.pop("subpatch_mapping")325                patch_metadata["timestamps"] = vd_metadata.timestamps326                patch_metadata["video_size"] = (vd_metadata.width, vd_metadata.height)327 328            video_grids = videos_inputs["video_grids"]329            # If user has not requested video metadata, pop it330            if "return_metadata" not in kwargs:331                video_metadata = videos_inputs.pop("video_metadata")332            else:333                video_metadata = videos_inputs["video_metadata"]334        else:335            videos_inputs = {}336            video_grids = None337 338        if not isinstance(text, list):339            text = [text]340        341        text = text.copy() # below lines change text in-place342 343        if image_grids is not None:344            index = 0345            for i in range(len(text)):346                num_images = text[i].count(self.image_placeholder_token)347                image_grids_i = image_grids[index:index+num_images]348                for image_grid in image_grids_i:349                    image_tokens = self.get_image_tokens(image_grid)350                    image_string = "".join(image_tokens)351                    text[i] = text[i].replace(self.image_placeholder_token, image_string, 1)352                index += num_images353        354        if video_grids is not None:355            index = 0356            for i in range(len(text)):357                num_videos = text[i].count(self.video_placeholder_token)358                assert num_videos in {0, 1}, "At most one video is supported for now"359                video_grids_i = video_grids[index:index+num_videos]360                metadata_i = video_metadata[index:index+num_videos]361                for video_grid, metadata in zip(video_grids_i, metadata_i):362                    video_string = self.get_video_string(363                        video_grid,364                        metadata.timestamps,365                    )366                    text[i] = text[i].replace(self.video_placeholder_token, video_string, 1)367                index += num_videos368 369        return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None)370        return_mm_token_type_ids = output_kwargs["text_kwargs"].pop("return_mm_token_type_ids", False)371        text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"])372 373        input_ids = text_inputs["input_ids"]374        attention_mask = text_inputs["attention_mask"]375 376        input_ids = np.array(input_ids)377        attention_mask = np.array(attention_mask)378        379        bos = self.tokenizer.bos_token_id or self.tokenizer.eos_token_id380        input_ids, attention_mask = self.insert_bos(381            input_ids, attention_mask, bos, self.tokenizer.pad_token_id382        )383 384        if return_mm_token_type_ids:385            image_tokens = np.array(self.image_token_ids).astype(input_ids.dtype)386            token_type_ids = np.any(input_ids[:, :, None] == image_tokens[None, None, :], axis=-1)387            text_inputs["token_type_ids"] = token_type_ids.tolist()388        389        text_inputs["input_ids"] = input_ids.tolist()390        text_inputs["attention_mask"] = attention_mask.tolist()391 392        features = BatchFeature(393            data={**text_inputs, **image_inputs, **videos_inputs},394            tensor_type=return_tensors,395        )396        if return_pointing_metadata:397            features["metadata"] = patch_metadata398        return features399 400    def post_process_image_text_to_text(401        self, generated_outputs, skip_special_tokens=True, clean_up_tokenization_spaces=False, **kwargs402    ):403        """404        Post-process the output of the model to decode the text.405 406        Args:407            generated_outputs (`torch.Tensor` or `np.ndarray`):408                The output of the model `generate` function. The output is expected to be a tensor of shape `(batch_size, sequence_length)`409                or `(sequence_length,)`.410            skip_special_tokens (`bool`, *optional*, defaults to `True`):411                Whether or not to remove special tokens in the output. Argument passed to the tokenizer's `batch_decode` method.412            clean_up_tokenization_spaces (`bool`, *optional*, defaults to `False`):413                Whether or not to clean up the tokenization spaces. Argument passed to the tokenizer's `batch_decode` method.414            **kwargs:415                Additional arguments to be passed to the tokenizer's `batch_decode method`.416 417        Returns:418            `list[str]`: The decoded text.419        """420        return self.tokenizer.batch_decode(421            generated_outputs,422            skip_special_tokens=skip_special_tokens,423            clean_up_tokenization_spaces=clean_up_tokenization_spaces,424            **kwargs,425        )426 427 428Molmo2Processor.register_for_auto_class()