from typing import Any, Optional from transformers.configuration_utils import PretrainedConfig from transformers.models.qwen3 import Qwen3Config from transformers import Qwen2_5_VLProcessor, AutoProcessor, AutoConfig from transformers.models.qwen2_5_vl.processing_qwen2_5_vl import Qwen2_5_VLProcessorKwargs, ImageInput, TextInput, PreTokenizedInput, VideoInput, BatchFeature, Unpack, Union, np class MonkeyOCRv2VisionConfig(PretrainedConfig): model_type: str = "monkeyocr_vit" def __init__( self, embed_dim: int = 1536, # vision encoder embed size hidden_size: int = 1536, # after merger hidden size intermediate_size: int = 4224, num_hidden_layers: int = 42, num_attention_heads: int = 12, num_channels: int = 3, patch_size: int = 14, spatial_merge_size: int = 2, temporal_patch_size: int = 1, rms_norm_eps: float = 1e-5, use_bias: bool = False, attn_implementation="flash_attention_2", # "eager","sdpa","flash_attention_2" initializer_range=0.02, init_merger_std=0.02, is_causal=False, # ve causal forward post_norm=True, gradient_checkpointing=False, **kwargs: Any, ): super().__init__(**kwargs) self.embed_dim = embed_dim self.hidden_size = hidden_size self.intermediate_size = intermediate_size self.num_hidden_layers = num_hidden_layers self.num_attention_heads = num_attention_heads self.num_channels = num_channels self.patch_size = patch_size self.spatial_merge_size = spatial_merge_size self.temporal_patch_size = temporal_patch_size self.rms_norm_eps = rms_norm_eps self.use_bias = use_bias self.attn_implementation = attn_implementation self.initializer_range = initializer_range self.init_merger_std = init_merger_std self.is_causal = is_causal self.post_norm = post_norm self.gradient_checkpointing = gradient_checkpointing class MonkeyOCRv2Config(Qwen3Config): model_type = "monkeyocrv2" def __init__(self, image_token_id = 151655, video_token_id = 151656, vision_config: Optional[dict] = None, *args, **kwargs): super().__init__(*args, **kwargs) self.image_token_id = image_token_id self.video_token_id = video_token_id self.vision_config = MonkeyOCRv2VisionConfig(**(vision_config or {})) def save_pretrained(self, save_directory, **kwargs): self._auto_class = None super().save_pretrained(save_directory, **kwargs) class MonkeyOCRv2Processor(Qwen2_5_VLProcessor): attributes = ["image_processor", "tokenizer"] def __init__(self, image_processor=None, tokenizer=None, chat_template=None, **kwargs): super().__init__(image_processor, tokenizer, chat_template=chat_template) self.image_token = "<|image_pad|>" if not hasattr(tokenizer, "image_token") else tokenizer.image_token self.image_token_id = 151655 if not hasattr(tokenizer, "image_token_id") else tokenizer.image_token_id def __call__( self, images: Optional[ImageInput] = None, text: Union[TextInput, PreTokenizedInput, list[TextInput], list[PreTokenizedInput]] = None, videos: Optional[VideoInput] = None, **kwargs: Unpack[Qwen2_5_VLProcessorKwargs], ) -> BatchFeature: """ Main method to prepare for the model one or several sequences(s) and image(s). This method forwards the `text` and `kwargs` arguments to Qwen2TokenizerFast's [`~Qwen2TokenizerFast.__call__`] if `text` is not `None` to encode the text. To prepare the vision inputs, this method forwards the `vision_infos` and `kwargs` arguments to Qwen2VLImageProcessor's [`~Qwen2VLImageProcessor.__call__`] if `vision_infos` is not `None`. Args: images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `list[PIL.Image.Image]`, `list[np.ndarray]`, `list[torch.Tensor]`): The image or batch of images to be prepared. Each image can be a PIL image, NumPy array or PyTorch tensor. Both channels-first and channels-last formats are supported. text (`str`, `list[str]`, `list[list[str]]`): The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set `is_split_into_words=True` (to lift the ambiguity with a batch of sequences). videos (`np.ndarray`, `torch.Tensor`, `list[np.ndarray]`, `list[torch.Tensor]`): The image or batch of videos to be prepared. Each video can be a 4D NumPy array or PyTorch tensor, or a nested list of 3D frames. Both channels-first and channels-last formats are supported. return_tensors (`str` or [`~utils.TensorType`], *optional*): If set, will return tensors of a particular framework. Acceptable values are: - `'tf'`: Return TensorFlow `tf.constant` objects. - `'pt'`: Return PyTorch `torch.Tensor` objects. - `'np'`: Return NumPy `np.ndarray` objects. - `'jax'`: Return JAX `jnp.ndarray` objects. Returns: [`BatchFeature`]: A [`BatchFeature`] with the following fields: - **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`. - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not `None`). - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`. - **pixel_values_videos** -- Pixel values of videos to be fed to a model. Returned when `videos` is not `None`. - **image_grid_thw** -- List of image 3D grid in LLM. Returned when `images` is not `None`. - **video_grid_thw** -- List of video 3D grid in LLM. Returned when `videos` is not `None`. - **second_per_grid_ts** -- List of video seconds per time grid. Returned when `videos` is not `None`. """ output_kwargs = self._merge_kwargs( Qwen2_5_VLProcessorKwargs, tokenizer_init_kwargs=self.tokenizer.init_kwargs, **kwargs, ) image_inputs = videos_inputs = {} if images is not None: image_inputs = self.image_processor(images=images, **output_kwargs["images_kwargs"]) image_grid_thw = image_inputs["image_grid_thw"] if videos is not None: fps = output_kwargs["videos_kwargs"].get("fps", 2.0) videos_inputs = self.video_processor(videos=videos, **output_kwargs["videos_kwargs"]) video_grid_thw = videos_inputs["video_grid_thw"] if isinstance(fps, (int, float)): second_per_grid_ts = [self.video_processor.temporal_patch_size / fps] * len(video_grid_thw) elif hasattr(fps, "__len__") and len(fps) == len(video_grid_thw): second_per_grid_ts = [self.video_processor.temporal_patch_size / tmp for tmp in fps] else: raise ValueError( f"The length of fps ({len(fps) if hasattr(fps, '__len__') else fps}) must be equal to the length of video_grid_thw ({len(video_grid_thw)}) or fps should be a single number." ) videos_inputs.update({"second_per_grid_ts": second_per_grid_ts}) if not isinstance(text, list): text = [text] text = text.copy() # below lines change text in-place if images is not None: merge_length = 1 #self.image_processor.merge_size**2 index = 0 for i in range(len(text)): while self.image_token in text[i]: num_image_tokens = image_grid_thw[index].prod() // merge_length text[i] = text[i].replace(self.image_token, "<|placeholder|>" * num_image_tokens, 1) index += 1 text[i] = text[i].replace("<|placeholder|>", self.image_token) if videos is not None: merge_length = self.video_processor.merge_size**2 index = 0 for i in range(len(text)): while self.video_token in text[i]: num_video_tokens = video_grid_thw[index].prod() // merge_length text[i] = text[i].replace(self.video_token, "<|placeholder|>" * num_video_tokens, 1) index += 1 text[i] = text[i].replace("<|placeholder|>", self.video_token) return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None) return_mm_token_type_ids = output_kwargs["text_kwargs"].pop("return_mm_token_type_ids", None) text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"]) self._check_special_mm_tokens(text, text_inputs, modalities=["image", "video"]) if return_mm_token_type_ids: array_ids = np.array(text_inputs["input_ids"]) mm_token_type_ids = np.zeros_like(text_inputs["input_ids"]) mm_token_type_ids[array_ids == self.image_token_id] = 1 text_inputs["mm_token_type_ids"] = mm_token_type_ids.tolist() return BatchFeature(data={**text_inputs, **image_inputs, **videos_inputs}, tensor_type=return_tensors) AutoProcessor.register("monkeyocrv2", MonkeyOCRv2Processor) AutoConfig.register("monkeyocrv2", MonkeyOCRv2Config) __all__ = ["MonkeyOCRv2Config", "MonkeyOCRv2VisionConfig", "MonkeyOCRv2Processor"]