pcuenq/nvidia-nano-clone
016
1from typing import Optional, Union, List2 3import numpy as np4 5from transformers.feature_extraction_utils import BatchFeature6from transformers.image_utils import ImageInput7from transformers.processing_utils import ImagesKwargs, MultiModalData, ProcessingKwargs, ProcessorMixin, Unpack, VideosKwargs8from transformers.tokenization_utils_base import PreTokenizedInput, TextInput9from transformers.video_utils import VideoInput10 11 12class NemotronNanoVLV2ImagesKwargs(ImagesKwargs):13 min_pixels: Optional[int]14 max_pixels: Optional[int]15 patch_size: Optional[int]16 temporal_patch_size: Optional[int]17 merge_size: Optional[int]18 19 20class NemotronNanoVLV2ProcessorKwargs(ProcessingKwargs, total=False):21 images_kwargs: NemotronNanoVLV2ImagesKwargs22 videos_kwargs: VideosKwargs23 _defaults = {24 "text_kwargs": {25 "padding": False,26 },27 }28 29 30class NemotronNanoVLV2Processor(ProcessorMixin):31 r"""32 Constructs a Nemotron Nano VL V2 processor which wraps an image processor and a tokenizer into a single processor.33 [`NemotronNanoVLV2Processor`] offers all the functionalities of the image processor and tokenizer. See the34 [`~NemotronNanoVLV2Processor.__call__`] and [`~NemotronNanoVLV2Processor.decode`] for more information.35 Args:36 image_processor ([`AutoImageProcessor`], *optional*):37 The image processor is a required input.38 tokenizer ([`AutoTokenizer`], *optional*):39 The tokenizer is a required input.40 chat_template (`str`, *optional*): A Jinja template which will be used to convert lists of messages41 in a chat into a tokenizable string.42 """43 44 attributes = ["image_processor", "tokenizer"]45 46 image_processor_class = "AutoImageProcessor"47 video_processor_class = "AutoVideoProcessor"48 tokenizer_class = ("AutoTokenizer")49 50 def __init__(self, image_processor=None, tokenizer=None, chat_template=None, **kwargs):51 self.image_token = "<image>" if not hasattr(tokenizer, "image_token") else tokenizer.image_token52 self.video_token = "<video>" if not hasattr(tokenizer, "video_token") else tokenizer.video_token53 self.image_start_token = "<img>" if not hasattr(tokenizer, "image_start_token") else tokenizer.image_start_token54 self.image_end_token = "</img>" if not hasattr(tokenizer, "image_end_token") else tokenizer.image_end_token55 self.image_token_id = (56 tokenizer.image_token_id57 if getattr(tokenizer, "image_token_id", None)58 else tokenizer.convert_tokens_to_ids(self.image_token)59 )60 self.video_token_id = (61 tokenizer.video_token_id62 if getattr(tokenizer, "video_token_id", None)63 else tokenizer.convert_tokens_to_ids(self.video_token)64 )65 super().__init__(image_processor, tokenizer, chat_template=chat_template)66 67 def __call__(68 self,69 images: ImageInput = None,70 text: Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]] = None,71 videos: VideoInput = None,72 **kwargs: Unpack[NemotronNanoVLV2ProcessorKwargs],73 ) -> BatchFeature:74 """75 Main method to prepare multimodal inputs (text, images, videos) for the model. This method processes text by 76 replacing image/video tokens with appropriate placeholder sequences, processes images and videos through the 77 image processor, and tokenizes the final text.78 79 The method performs the following key operations:80 1. Processes images using the image processor to get pixel values and patch counts81 2. Processes videos using the image processor with max_num_tiles=1 to get video pixel values 82 3. Replaces `<image>` tokens in text with `<img>` + image tokens + `</img>` sequences83 4. Replaces `<video>` tokens in text with frame-by-frame descriptions including timestamps (if metadata provided)84 5. Tokenizes the processed text and combines all outputs85 86 Args:87 images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `List[PIL.Image.Image]`, `List[np.ndarray]`, `List[torch.Tensor]`, *optional*):88 The image or batch of images to be prepared. Each image can be a PIL image, NumPy array or PyTorch89 tensor. Both channels-first and channels-last formats are supported.90 text (`str`, `List[str]`, *optional*):91 The sequence or batch of sequences to be encoded. Each sequence should be a string. The text can contain92 special tokens `<image>` and `<video>` that will be replaced with appropriate token sequences.93 videos (`np.ndarray`, `torch.Tensor`, `List[np.ndarray]`, `List[torch.Tensor]`, *optional*):94 The video or batch of videos to be prepared. Each video should be a 4D NumPy array or PyTorch95 tensor with shape (num_frames, channels, height, width). Both channels-first and channels-last formats 96 are supported. Note: Currently only supports batch size of 1 for videos.97 images_kwargs (`Dict`, *optional*):98 Additional keyword arguments for image processing, including:99 - `min_pixels` (`int`, *optional*): Minimum number of pixels for image processing100 - `max_pixels` (`int`, *optional*): Maximum number of pixels for image processing 101 - `patch_size` (`int`, *optional*): Size of patches for image processing102 - `temporal_patch_size` (`int`, *optional*): Size of temporal patches103 - `merge_size` (`int`, *optional*): Size for merging patches104 videos_kwargs (`Dict`, *optional*):105 Additional keyword arguments for video processing, including:106 - `video_metadata` (`VideoMetadata`, *optional*): Metadata containing fps information for timestamp calculation107 text_kwargs (`Dict`, *optional*):108 Additional keyword arguments for text tokenization, including:109 - `return_tensors` (`str` or [`~utils.TensorType`], *optional*): Framework for returned tensors ('tf', 'pt', 'np', 'jax')110 - `padding` (`bool`, *optional*): Whether to pad sequences (defaults to False)111 112 Returns:113 [`BatchFeature`]: A [`BatchFeature`] with the following fields:114 115 - **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`.116 - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when117 `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not118 `None`).119 - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`.120 - **num_patches** -- Number of patches per image. Returned when `images` is not `None`.121 - **pixel_values_videos** -- Pixel values of videos to be fed to a model. Returned when `videos` is not `None`.122 123 Raises:124 AssertionError: If videos are provided with batch size > 1 (not currently supported).125 126 Note:127 - Image tokens `<image>` in text are replaced with `<img>` + repeated image tokens + `</img>`128 - Video tokens `<video>` in text are replaced with frame-by-frame descriptions129 - When video metadata with fps is provided, frame descriptions include timestamps130 - Videos are processed with max_num_tiles=1 regardless of the images setting131 """132 output_kwargs = self._merge_kwargs(133 NemotronNanoVLV2ProcessorKwargs,134 tokenizer_init_kwargs=self.tokenizer.init_kwargs,135 **kwargs,136 )137 image_inputs = videos_inputs = {}138 if images is not None:139 image_inputs = self.image_processor(images=images, **output_kwargs["images_kwargs"])140 image_num_patches = image_inputs["num_patches"]141 142 if videos is not None:143 orig_tiles = self.image_processor.max_num_tiles144 self.image_processor.max_num_tiles = 1145 videos_inputs = self.image_processor(images=videos, **output_kwargs["images_kwargs"])146 self.image_processor.max_num_tiles = orig_tiles147 video_num_patches = [sum(videos_inputs["num_patches"])]148 videos_inputs["pixel_values_videos"] = videos_inputs["pixel_values"]149 del videos_inputs["pixel_values"]150 151 if not isinstance(text, list):152 text = [text]153 154 text = text.copy() # below lines change text in-place155 if images is not None:156 index = 0157 for i in range(len(text)):158 while self.image_token in text[i]:159 text[i] = text[i].replace(self.image_token, self.image_start_token + "<|placeholder|>" * image_num_patches[index] * self.image_processor.num_image_token + self.image_end_token, 1)160 index += 1161 text[i] = text[i].replace("<|placeholder|>", self.image_token)162 if videos is not None:163 assert len(text) == 1, "Video is not supported for batch size > 1"164 video_metadata = output_kwargs.get("videos_kwargs", {}).get("video_metadata", None)165 i = 0166 index = 0167 if self.video_token in text[i]:168 each_frame = self.image_start_token + "<|placeholder|>" * self.image_processor.num_image_token + self.image_end_token169 video_prompt = "This is a video:\n"170 for j in range(video_num_patches[index]):171 if video_metadata is not None and video_metadata.fps is not None:172 timestamp = j / video_metadata.fps173 video_prompt += f"Frame {j+1} sampled at {timestamp:.2f} seconds: {each_frame}\n"174 else:175 # Fallback to original format without timestamps176 video_prompt += f"Frame {j+1}: {each_frame}\n"177 178 text[i] = text[i].replace(self.video_token, video_prompt, 1)179 text[i] = text[i].replace("<|placeholder|>", self.video_token)180 181 return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None)182 text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"])183 return BatchFeature(data={**text_inputs, **image_inputs, **videos_inputs}, tensor_type=return_tensors)184 185 def _get_num_multimodal_tokens(self, image_sizes=None, video_sizes=None, **kwargs):186 """187 Computes the number of placeholder tokens needed for multimodal inputs with the given sizes.188 Args:189 image_sizes (`list[list[int]]`, *optional*):190 The input sizes formatted as (height, width) per each image.191 video_sizes (`list[list[int]]`, *optional*):192 The input sizes formatted as (num_frames, height, width) per each video.193 Returns:194 `MultiModalData`: A `MultiModalData` object holding number of tokens per each of the provided195 input modalities, along with other useful data.196 """197 198 vision_data = {}199 if image_sizes is not None:200 images_kwargs = NemotronNanoVLV2ProcessorKwargs._defaults.get("images_kwargs", {})201 images_kwargs.update(kwargs)202 merge_size = images_kwargs.get("merge_size", None) or self.image_processor.merge_size203 204 num_image_patches = [205 self.image_processor.get_number_of_image_patches(*image_size, images_kwargs)206 for image_size in image_sizes207 ]208 num_image_tokens = [(num_patches // merge_size**2) for num_patches in num_image_patches]209 vision_data.update({"num_image_tokens": num_image_tokens, "num_image_patches": num_image_patches})210 return MultiModalData(**vision_data)211 212 def batch_decode(self, *args, **kwargs):213 """214 This method forwards all its arguments to the tokenizer's [`~PreTrainedTokenizer.batch_decode`]. Please215 refer to the docstring of this method for more information.216 """217 return self.tokenizer.batch_decode(*args, **kwargs)218 219 def decode(self, *args, **kwargs):220 """221 This method forwards all its arguments to the tokenizer's [`~PreTrainedTokenizer.decode`]. Please refer to222 the docstring of this method for more information.223 """224 return self.tokenizer.decode(*args, **kwargs)225 226 def post_process_image_text_to_text(227 self, generated_outputs, skip_special_tokens=True, clean_up_tokenization_spaces=False, **kwargs228 ):229 """230 Post-process the output of the model to decode the text.231 232 Args:233 generated_outputs (`torch.Tensor` or `np.ndarray`):234 The output of the model `generate` function. The output is expected to be a tensor of shape `(batch_size, sequence_length)`235 or `(sequence_length,)`.236 skip_special_tokens (`bool`, *optional*, defaults to `True`):237 Whether or not to remove special tokens in the output. Argument passed to the tokenizer's `batch_decode` method.238 clean_up_tokenization_spaces (`bool`, *optional*, defaults to `False`):239 Whether or not to clean up the tokenization spaces. Argument passed to the tokenizer's `batch_decode` method.240 **kwargs:241 Additional arguments to be passed to the tokenizer's `batch_decode method`.242 243 Returns:244 `list[str]`: The decoded text.245 """246 return self.tokenizer.batch_decode(247 generated_outputs,248 skip_special_tokens=skip_special_tokens,249 clean_up_tokenization_spaces=clean_up_tokenization_spaces,250 **kwargs,251 )252 253 @property254 def model_input_names(self):255 tokenizer_input_names = self.tokenizer.model_input_names256 image_processor_input_names = self.image_processor.model_input_names257 names_from_processor = list(dict.fromkeys(tokenizer_input_names + image_processor_input_names))258 return names_from_processor + ["second_per_grid_ts"]259 260 261__all__ = ["NemotronNanoVLV2Processor"]262