|
|
from typing import Optional, Union, List |
|
|
|
|
|
import numpy as np |
|
|
|
|
|
from transformers.feature_extraction_utils import BatchFeature |
|
|
from transformers.image_utils import ImageInput |
|
|
from transformers.processing_utils import ImagesKwargs, MultiModalData, ProcessingKwargs, ProcessorMixin, Unpack, VideosKwargs |
|
|
from transformers.tokenization_utils_base import PreTokenizedInput, TextInput |
|
|
from transformers.video_utils import VideoInput |
|
|
|
|
|
|
|
|
class NemotronNanoVLV2ImagesKwargs(ImagesKwargs): |
|
|
min_pixels: Optional[int] |
|
|
max_pixels: Optional[int] |
|
|
patch_size: Optional[int] |
|
|
temporal_patch_size: Optional[int] |
|
|
merge_size: Optional[int] |
|
|
|
|
|
|
|
|
class NemotronNanoVLV2ProcessorKwargs(ProcessingKwargs, total=False): |
|
|
images_kwargs: NemotronNanoVLV2ImagesKwargs |
|
|
videos_kwargs: VideosKwargs |
|
|
_defaults = { |
|
|
"text_kwargs": { |
|
|
"padding": False, |
|
|
}, |
|
|
} |
|
|
|
|
|
|
|
|
class NemotronNanoVLV2Processor(ProcessorMixin): |
|
|
r""" |
|
|
Constructs a Nemotron Nano VL V2 processor which wraps an image processor and a tokenizer into a single processor. |
|
|
[`NemotronNanoVLV2Processor`] offers all the functionalities of the image processor and tokenizer. See the |
|
|
[`~NemotronNanoVLV2Processor.__call__`] and [`~NemotronNanoVLV2Processor.decode`] for more information. |
|
|
Args: |
|
|
image_processor ([`AutoImageProcessor`], *optional*): |
|
|
The image processor is a required input. |
|
|
tokenizer ([`AutoTokenizer`], *optional*): |
|
|
The tokenizer is a required input. |
|
|
chat_template (`str`, *optional*): A Jinja template which will be used to convert lists of messages |
|
|
in a chat into a tokenizable string. |
|
|
""" |
|
|
|
|
|
attributes = ["image_processor", "tokenizer"] |
|
|
|
|
|
image_processor_class = "AutoImageProcessor" |
|
|
video_processor_class = "AutoVideoProcessor" |
|
|
tokenizer_class = ("AutoTokenizer") |
|
|
|
|
|
def __init__(self, image_processor=None, tokenizer=None, chat_template=None, **kwargs): |
|
|
self.image_token = "<image>" if not hasattr(tokenizer, "image_token") else tokenizer.image_token |
|
|
self.video_token = "<video>" if not hasattr(tokenizer, "video_token") else tokenizer.video_token |
|
|
self.image_start_token = "<img>" if not hasattr(tokenizer, "image_start_token") else tokenizer.image_start_token |
|
|
self.image_end_token = "</img>" if not hasattr(tokenizer, "image_end_token") else tokenizer.image_end_token |
|
|
self.image_token_id = ( |
|
|
tokenizer.image_token_id |
|
|
if getattr(tokenizer, "image_token_id", None) |
|
|
else tokenizer.convert_tokens_to_ids(self.image_token) |
|
|
) |
|
|
self.video_token_id = ( |
|
|
tokenizer.video_token_id |
|
|
if getattr(tokenizer, "video_token_id", None) |
|
|
else tokenizer.convert_tokens_to_ids(self.video_token) |
|
|
) |
|
|
super().__init__(image_processor, tokenizer, chat_template=chat_template) |
|
|
|
|
|
def __call__( |
|
|
self, |
|
|
images: ImageInput = None, |
|
|
text: Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]] = None, |
|
|
videos: VideoInput = None, |
|
|
**kwargs: Unpack[NemotronNanoVLV2ProcessorKwargs], |
|
|
) -> BatchFeature: |
|
|
""" |
|
|
Main method to prepare multimodal inputs (text, images, videos) for the model. This method processes text by |
|
|
replacing image/video tokens with appropriate placeholder sequences, processes images and videos through the |
|
|
image processor, and tokenizes the final text. |
|
|
|
|
|
The method performs the following key operations: |
|
|
1. Processes images using the image processor to get pixel values and patch counts |
|
|
2. Processes videos using the image processor with max_num_tiles=1 to get video pixel values |
|
|
3. Replaces `<image>` tokens in text with `<img>` + image tokens + `</img>` sequences |
|
|
4. Replaces `<video>` tokens in text with frame-by-frame descriptions including timestamps (if metadata provided) |
|
|
5. Tokenizes the processed text and combines all outputs |
|
|
|
|
|
Args: |
|
|
images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `List[PIL.Image.Image]`, `List[np.ndarray]`, `List[torch.Tensor]`, *optional*): |
|
|
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]`, *optional*): |
|
|
The sequence or batch of sequences to be encoded. Each sequence should be a string. The text can contain |
|
|
special tokens `<image>` and `<video>` that will be replaced with appropriate token sequences. |
|
|
videos (`np.ndarray`, `torch.Tensor`, `List[np.ndarray]`, `List[torch.Tensor]`, *optional*): |
|
|
The video or batch of videos to be prepared. Each video should be a 4D NumPy array or PyTorch |
|
|
tensor with shape (num_frames, channels, height, width). Both channels-first and channels-last formats |
|
|
are supported. Note: Currently only supports batch size of 1 for videos. |
|
|
images_kwargs (`Dict`, *optional*): |
|
|
Additional keyword arguments for image processing, including: |
|
|
- `min_pixels` (`int`, *optional*): Minimum number of pixels for image processing |
|
|
- `max_pixels` (`int`, *optional*): Maximum number of pixels for image processing |
|
|
- `patch_size` (`int`, *optional*): Size of patches for image processing |
|
|
- `temporal_patch_size` (`int`, *optional*): Size of temporal patches |
|
|
- `merge_size` (`int`, *optional*): Size for merging patches |
|
|
videos_kwargs (`Dict`, *optional*): |
|
|
Additional keyword arguments for video processing, including: |
|
|
- `video_metadata` (`VideoMetadata`, *optional*): Metadata containing fps information for timestamp calculation |
|
|
text_kwargs (`Dict`, *optional*): |
|
|
Additional keyword arguments for text tokenization, including: |
|
|
- `return_tensors` (`str` or [`~utils.TensorType`], *optional*): Framework for returned tensors ('tf', 'pt', 'np', 'jax') |
|
|
- `padding` (`bool`, *optional*): Whether to pad sequences (defaults to False) |
|
|
|
|
|
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`. |
|
|
- **num_patches** -- Number of patches per image. 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`. |
|
|
|
|
|
Raises: |
|
|
AssertionError: If videos are provided with batch size > 1 (not currently supported). |
|
|
|
|
|
Note: |
|
|
- Image tokens `<image>` in text are replaced with `<img>` + repeated image tokens + `</img>` |
|
|
- Video tokens `<video>` in text are replaced with frame-by-frame descriptions |
|
|
- When video metadata with fps is provided, frame descriptions include timestamps |
|
|
- Videos are processed with max_num_tiles=1 regardless of the images setting |
|
|
""" |
|
|
output_kwargs = self._merge_kwargs( |
|
|
NemotronNanoVLV2ProcessorKwargs, |
|
|
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_num_patches = image_inputs["num_patches"] |
|
|
|
|
|
if videos is not None: |
|
|
orig_tiles = self.image_processor.max_num_tiles |
|
|
self.image_processor.max_num_tiles = 1 |
|
|
videos_inputs = self.image_processor(images=videos, **output_kwargs["images_kwargs"]) |
|
|
self.image_processor.max_num_tiles = orig_tiles |
|
|
video_num_patches = [sum(videos_inputs["num_patches"])] |
|
|
videos_inputs["pixel_values_videos"] = videos_inputs["pixel_values"] |
|
|
del videos_inputs["pixel_values"] |
|
|
|
|
|
if not isinstance(text, list): |
|
|
text = [text] |
|
|
|
|
|
text = text.copy() |
|
|
if images is not None: |
|
|
index = 0 |
|
|
for i in range(len(text)): |
|
|
while self.image_token in text[i]: |
|
|
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) |
|
|
index += 1 |
|
|
text[i] = text[i].replace("<|placeholder|>", self.image_token) |
|
|
if videos is not None: |
|
|
assert len(text) == 1, "Video is not supported for batch size > 1" |
|
|
video_metadata = output_kwargs.get("videos_kwargs", {}).get("video_metadata", None) |
|
|
i = 0 |
|
|
index = 0 |
|
|
if self.video_token in text[i]: |
|
|
each_frame = self.image_start_token + "<|placeholder|>" * self.image_processor.num_image_token + self.image_end_token |
|
|
video_prompt = "This is a video:\n" |
|
|
for j in range(video_num_patches[index]): |
|
|
if video_metadata is not None and video_metadata.fps is not None: |
|
|
timestamp = j / video_metadata.fps |
|
|
video_prompt += f"Frame {j+1} sampled at {timestamp:.2f} seconds: {each_frame}\n" |
|
|
else: |
|
|
|
|
|
video_prompt += f"Frame {j+1}: {each_frame}\n" |
|
|
|
|
|
text[i] = text[i].replace(self.video_token, video_prompt, 1) |
|
|
text[i] = text[i].replace("<|placeholder|>", self.video_token) |
|
|
|
|
|
return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None) |
|
|
text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"]) |
|
|
return BatchFeature(data={**text_inputs, **image_inputs, **videos_inputs}, tensor_type=return_tensors) |
|
|
|
|
|
def _get_num_multimodal_tokens(self, image_sizes=None, video_sizes=None, **kwargs): |
|
|
""" |
|
|
Computes the number of placeholder tokens needed for multimodal inputs with the given sizes. |
|
|
Args: |
|
|
image_sizes (`list[list[int]]`, *optional*): |
|
|
The input sizes formatted as (height, width) per each image. |
|
|
video_sizes (`list[list[int]]`, *optional*): |
|
|
The input sizes formatted as (num_frames, height, width) per each video. |
|
|
Returns: |
|
|
`MultiModalData`: A `MultiModalData` object holding number of tokens per each of the provided |
|
|
input modalities, along with other useful data. |
|
|
""" |
|
|
|
|
|
vision_data = {} |
|
|
if image_sizes is not None: |
|
|
images_kwargs = NemotronNanoVLV2ProcessorKwargs._defaults.get("images_kwargs", {}) |
|
|
images_kwargs.update(kwargs) |
|
|
merge_size = images_kwargs.get("merge_size", None) or self.image_processor.merge_size |
|
|
|
|
|
num_image_patches = [ |
|
|
self.image_processor.get_number_of_image_patches(*image_size, images_kwargs) |
|
|
for image_size in image_sizes |
|
|
] |
|
|
num_image_tokens = [(num_patches // merge_size**2) for num_patches in num_image_patches] |
|
|
vision_data.update({"num_image_tokens": num_image_tokens, "num_image_patches": num_image_patches}) |
|
|
return MultiModalData(**vision_data) |
|
|
|
|
|
def batch_decode(self, *args, **kwargs): |
|
|
""" |
|
|
This method forwards all its arguments to the tokenizer's [`~PreTrainedTokenizer.batch_decode`]. Please |
|
|
refer to the docstring of this method for more information. |
|
|
""" |
|
|
return self.tokenizer.batch_decode(*args, **kwargs) |
|
|
|
|
|
def decode(self, *args, **kwargs): |
|
|
""" |
|
|
This method forwards all its arguments to the tokenizer's [`~PreTrainedTokenizer.decode`]. Please refer to |
|
|
the docstring of this method for more information. |
|
|
""" |
|
|
return self.tokenizer.decode(*args, **kwargs) |
|
|
|
|
|
def post_process_image_text_to_text( |
|
|
self, generated_outputs, skip_special_tokens=True, clean_up_tokenization_spaces=False, **kwargs |
|
|
): |
|
|
""" |
|
|
Post-process the output of the model to decode the text. |
|
|
|
|
|
Args: |
|
|
generated_outputs (`torch.Tensor` or `np.ndarray`): |
|
|
The output of the model `generate` function. The output is expected to be a tensor of shape `(batch_size, sequence_length)` |
|
|
or `(sequence_length,)`. |
|
|
skip_special_tokens (`bool`, *optional*, defaults to `True`): |
|
|
Whether or not to remove special tokens in the output. Argument passed to the tokenizer's `batch_decode` method. |
|
|
clean_up_tokenization_spaces (`bool`, *optional*, defaults to `False`): |
|
|
Whether or not to clean up the tokenization spaces. Argument passed to the tokenizer's `batch_decode` method. |
|
|
**kwargs: |
|
|
Additional arguments to be passed to the tokenizer's `batch_decode method`. |
|
|
|
|
|
Returns: |
|
|
`list[str]`: The decoded text. |
|
|
""" |
|
|
return self.tokenizer.batch_decode( |
|
|
generated_outputs, |
|
|
skip_special_tokens=skip_special_tokens, |
|
|
clean_up_tokenization_spaces=clean_up_tokenization_spaces, |
|
|
**kwargs, |
|
|
) |
|
|
|
|
|
@property |
|
|
def model_input_names(self): |
|
|
tokenizer_input_names = self.tokenizer.model_input_names |
|
|
image_processor_input_names = self.image_processor.model_input_names |
|
|
names_from_processor = list(dict.fromkeys(tokenizer_input_names + image_processor_input_names)) |
|
|
return names_from_processor + ["second_per_grid_ts"] |
|
|
|
|
|
|
|
|
__all__ = ["NemotronNanoVLV2Processor"] |
|
|
|
