Pixtral: vectorize patch embeddings and enable tests (#35122)

* initial POC

* - batch mix feature

* fix tests

* fix tests

* make style

* do not skip and instead fix tests

* update

* return back the test

* correct text with the correct ckpt
This commit is contained in:
Raushan Turganbay 2025-01-30 12:40:18 +01:00 committed by GitHub
parent 8bc4c89ee9
commit 9725e5be2f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 422 additions and 545 deletions

View file

@ -280,6 +280,7 @@ class LlavaForConditionalGeneration(LlavaPreTrainedModel, GenerationMixin):
pixel_values: torch.FloatTensor,
vision_feature_layer: Union[int, List[int]],
vision_feature_select_strategy: str,
**kwargs,
):
"""
Obtains image last hidden states from the vision tower and apply multimodal projection.
@ -300,8 +301,9 @@ class LlavaForConditionalGeneration(LlavaPreTrainedModel, GenerationMixin):
if vision_feature_select_strategy not in ["default", "full"]:
raise ValueError(f"Unexpected select feature strategy: {self.config.vision_feature_select_strategy}")
kwargs = {k: v for k, v in kwargs.items() if v is not None}
# this is not memory efficient at all (output_hidden_states=True) will save all the hidden states.
image_outputs = self.vision_tower(pixel_values, output_hidden_states=True)
image_outputs = self.vision_tower(pixel_values, output_hidden_states=True, **kwargs)
# If we have one vision feature layer, return the corresponding hidden states,
# otherwise, select the hidden states of each feature layer and concatenate them
@ -422,6 +424,7 @@ class LlavaForConditionalGeneration(LlavaPreTrainedModel, GenerationMixin):
return_dict: Optional[bool] = None,
cache_position: Optional[torch.LongTensor] = None,
logits_to_keep: Union[int, torch.Tensor] = 0,
image_sizes: torch.Tensor = None,
) -> Union[Tuple, LlavaCausalLMOutputWithPast]:
r"""
Args:
@ -492,6 +495,7 @@ class LlavaForConditionalGeneration(LlavaPreTrainedModel, GenerationMixin):
pixel_values=pixel_values,
vision_feature_layer=vision_feature_layer,
vision_feature_select_strategy=vision_feature_select_strategy,
image_sizes=image_sizes,
)
n_image_tokens = (input_ids == self.config.image_token_index).sum().item()

View file

@ -15,12 +15,13 @@
"""Image processor class for Pixtral."""
import math
from typing import Any, Callable, Dict, List, Optional, Tuple, Union
from typing import Dict, List, Optional, Tuple, Union
import numpy as np
from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import (
pad,
resize,
to_channel_dimension_format,
)
@ -31,13 +32,13 @@ from ...image_utils import (
get_image_size,
infer_channel_dimension_format,
is_scaled_image,
is_valid_image,
make_list_of_images,
to_numpy_array,
valid_images,
validate_kwargs,
validate_preprocess_arguments,
)
from ...utils import TensorType, is_torch_device, is_torch_dtype, is_vision_available, logging
from ...utils import TensorType, is_vision_available, logging
from ...utils.import_utils import requires_backends
@ -48,91 +49,6 @@ if is_vision_available():
import PIL
class BatchMixFeature(BatchFeature):
def to(self, *args, **kwargs) -> "BatchMixFeature":
"""
Send all values to device by calling `v.to(*args, **kwargs)` (PyTorch only). This should support casting in
different `dtypes` and sending the `BatchFeature` to a different `device`.
Args:
args (`Tuple`):
Will be passed to the `to(...)` function of the tensors.
kwargs (`Dict`, *optional*):
Will be passed to the `to(...)` function of the tensors.
Returns:
[`BatchFeature`]: The same instance after modification.
"""
def _recursive_to(obj, device, *args, **kwargs):
# Lists can be nested, so keep digging until we hit tensors
if isinstance(obj, list):
return [_recursive_to(o, device, *args, **kwargs) for o in obj]
# We cast only floating point tensors to avoid issues with tokenizers casting `LongTensor` to `FloatTensor`
elif isinstance(obj, torch.Tensor) and torch.is_floating_point(obj):
# cast and send to device
return obj.to(*args, **kwargs)
elif isinstance(obj, torch.Tensor) and device is not None:
# only send to device, don't cast
return obj.to(device=device)
else:
return obj
requires_backends(self, ["torch"])
import torch # noqa
device = kwargs.get("device")
# Check if the args are a device or a dtype
if device is None and len(args) > 0:
# device should be always the first argument
arg = args[0]
if is_torch_dtype(arg):
# The first argument is a dtype
pass
elif isinstance(arg, str) or is_torch_device(arg) or isinstance(arg, int):
device = arg
else:
# it's something else
raise ValueError(f"Attempting to cast a BatchFeature to type {str(arg)}. This is not supported.")
self.data = {k: _recursive_to(v, device, *args, **kwargs) for k, v in self.data.items()}
return self
# Copied from transformers.models.idefics2.image_processing_idefics2.make_list_of_images
def make_list_of_images(images: ImageInput) -> List[List[np.ndarray]]:
"""
Convert a single image or a list of images to a list of numpy arrays.
Args:
images (`ImageInput`):
A single image or a list of images.
Returns:
A list of numpy arrays.
"""
# If it's a single image, convert it to a list of lists
if is_valid_image(images):
images = [[images]]
# If it's a list of images, it's a single batch, so convert it to a list of lists
elif isinstance(images, (list, tuple)) and len(images) > 0 and is_valid_image(images[0]):
images = [images]
# If it's a list of batches, it's already in the right format
elif (
isinstance(images, (list, tuple))
and len(images) > 0
and isinstance(images[0], (list, tuple))
and len(images[0]) > 0
and is_valid_image(images[0][0])
):
pass
else:
raise ValueError(
"Invalid input type. Must be a single image, a list of images, or a list of batches of images."
)
return images
# Adapted from function in image_transforms.py to ensure any transparent pixels are converted to white.
def convert_to_rgb(image: ImageInput) -> ImageInput:
"""
@ -219,18 +135,6 @@ def get_resize_output_image_size(
return num_height_tokens * patch_height, num_width_tokens * patch_width
# Hack to get tensor conversion used in BatchFeature without batching the images
def _get_is_as_tensor_fns(tensor_type: Union[str, TensorType]) -> Tuple[Callable, Callable]:
return BatchFeature()._get_is_as_tensor_fns(tensor_type)
def convert_to_tensor(array, tensor_type: Union[str, TensorType]) -> Any:
is_tensor, as_tensor = _get_is_as_tensor_fns(tensor_type)
if is_tensor(array):
return array
return as_tensor(array)
class PixtralImageProcessor(BaseImageProcessor):
r"""
Constructs a Pixtral image processor.
@ -368,6 +272,49 @@ class PixtralImageProcessor(BaseImageProcessor):
**kwargs,
)
def _pad_for_batching(
self,
pixel_values: List[np.ndarray],
image_sizes: List[List[int]],
data_format: Optional[Union[str, ChannelDimension]] = None,
input_data_format: Optional[Union[str, ChannelDimension]] = None,
):
"""
Pads images on the `num_of_patches` dimension with zeros to form a batch of same number of patches.
Args:
pixel_values (`List[np.ndarray]`):
An array of pixel values of each images of shape (`batch_size`, `height`, `width`, `channels`)
image_sizes (`List[List[int]]`):
A list of sizes for each image in `pixel_values` in (height, width) format.
data_format (`str` or `ChannelDimension`, *optional*):
The channel dimension format for the output image. Can be one of:
- `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
- `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
If unset, will use same as the input image.
input_data_format (`str` or `ChannelDimension`, *optional*):
The channel dimension format for the input image. Can be one of:
- `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
- `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
If unset, will use the inferred format of the input image.
Returns:
List[`np.ndarray`]: The padded images.
"""
max_shape = (
max([size[0] for size in image_sizes]),
max([size[1] for size in image_sizes]),
)
pixel_values = [
pad(
image,
padding=((0, max_shape[0] - size[0]), (0, max_shape[1] - size[1])),
data_format=data_format,
input_data_format=input_data_format,
)
for image, size in zip(pixel_values, image_sizes)
]
return pixel_values
def preprocess(
self,
images: ImageInput,
@ -449,9 +396,9 @@ class PixtralImageProcessor(BaseImageProcessor):
validate_kwargs(captured_kwargs=kwargs.keys(), valid_processor_keys=self._valid_processor_keys)
images_list = make_list_of_images(images)
images = make_list_of_images(images)
if not valid_images(images_list[0]):
if not valid_images(images[0]):
raise ValueError(
"Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "
"torch.Tensor, tf.Tensor or jax.ndarray."
@ -469,12 +416,12 @@ class PixtralImageProcessor(BaseImageProcessor):
)
if do_convert_rgb:
images_list = [[convert_to_rgb(image) for image in images] for images in images_list]
images = [convert_to_rgb(image) for image in images]
# All transformations expect numpy arrays.
images_list = [[to_numpy_array(image) for image in images] for images in images_list]
images = [to_numpy_array(image) for image in images]
if do_rescale and is_scaled_image(images_list[0][0]):
if do_rescale and is_scaled_image(images[0]):
logger.warning_once(
"It looks like you are trying to rescale already rescaled images. If the input"
" images have pixel values between 0 and 1, set `do_rescale=False` to avoid rescaling them again."
@ -482,44 +429,43 @@ class PixtralImageProcessor(BaseImageProcessor):
if input_data_format is None:
# We assume that all images have the same channel dimension format.
input_data_format = infer_channel_dimension_format(images_list[0][0])
input_data_format = infer_channel_dimension_format(images[0])
batch_images = []
batch_image_sizes = []
for sample_images in images_list:
images = []
image_sizes = []
for image in sample_images:
if do_resize:
image = self.resize(
image=image,
size=size,
patch_size=patch_size,
resample=resample,
input_data_format=input_data_format,
)
for image in images:
if do_resize:
image = self.resize(
image=image,
size=size,
patch_size=patch_size,
resample=resample,
input_data_format=input_data_format,
)
if do_rescale:
image = self.rescale(image=image, scale=rescale_factor, input_data_format=input_data_format)
if do_rescale:
image = self.rescale(image=image, scale=rescale_factor, input_data_format=input_data_format)
if do_normalize:
image = self.normalize(
image=image, mean=image_mean, std=image_std, input_data_format=input_data_format
)
if do_normalize:
image = self.normalize(
image=image, mean=image_mean, std=image_std, input_data_format=input_data_format
)
images.append(image)
image_sizes.append(get_image_size(image, input_data_format))
batch_images.append(images)
batch_image_sizes.append(image_sizes)
image = to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format)
images_list = [
[to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format) for image in images]
for images in batch_images
]
batch_images.append(image)
batch_image_sizes.append(get_image_size(image, data_format))
# Convert to tensor type outside of BatchFeature to avoid batching the images of different sizes
images_list = [[convert_to_tensor(image, return_tensors) for image in images] for images in images_list]
return BatchMixFeature(data={"pixel_values": images_list, "image_sizes": batch_image_sizes}, tensor_type=None)
pixel_values = self._pad_for_batching(
pixel_values=batch_images,
image_sizes=batch_image_sizes,
input_data_format=data_format,
data_format=data_format,
)
return BatchFeature(
data={"pixel_values": pixel_values, "image_sizes": batch_image_sizes}, tensor_type=return_tensors
)
__all__ = ["PixtralImageProcessor"]

View file

@ -16,7 +16,7 @@
from typing import Dict, List, Optional, Union
from ...image_processing_utils import get_size_dict
from ...image_processing_utils import BatchFeature, get_size_dict
from ...image_processing_utils_fast import BaseImageProcessorFast
from ...image_utils import (
ChannelDimension,
@ -26,6 +26,7 @@ from ...image_utils import (
get_image_size,
get_image_type,
infer_channel_dimension_format,
make_list_of_images,
validate_fast_preprocess_arguments,
validate_kwargs,
)
@ -38,10 +39,8 @@ from ...utils import (
logging,
)
from .image_processing_pixtral import (
BatchMixFeature,
convert_to_rgb,
get_resize_output_image_size,
make_list_of_images,
)
@ -189,6 +188,36 @@ class PixtralImageProcessorFast(BaseImageProcessorFast):
**kwargs,
)
# Adapted from transformers.models.pixtral.image_processing_pixtral.PixtralImageProcessor._pad_for_batching
def _pad_for_batching(
self,
pixel_values: List[torch.Tensor],
image_sizes: List[List[int]],
):
"""
Pads images on the `num_of_patches` dimension with zeros to form a batch of same number of patches.
Args:
pixel_values (`List[torch.Tensor]`):
An array of pixel values of each images of shape (`batch_size`, `channels`, `height`, `width`)
image_sizes (`List[List[int]]`):
A list of sizes for each image in `pixel_values` in (height, width) format.
Returns:
List[`torch.Tensor`]: The padded images.
"""
max_shape = (
max([size[0] for size in image_sizes]),
max([size[1] for size in image_sizes]),
)
pixel_values = [
torch.nn.functional.pad(
image,
pad=(0, max_shape[1] - size[1], 0, max_shape[0] - size[0]),
)
for image, size in zip(pixel_values, image_sizes)
]
return torch.stack(pixel_values)
def preprocess(
self,
images: ImageInput,
@ -206,7 +235,7 @@ class PixtralImageProcessorFast(BaseImageProcessorFast):
data_format: Optional[ChannelDimension] = ChannelDimension.FIRST,
input_data_format: Optional[Union[str, ChannelDimension]] = None,
**kwargs,
) -> BatchMixFeature:
) -> BatchFeature:
"""
Preprocess an image or batch of images.
@ -271,8 +300,8 @@ class PixtralImageProcessorFast(BaseImageProcessorFast):
validate_kwargs(captured_kwargs=kwargs.keys(), valid_processor_keys=self._valid_processor_keys)
images_list = make_list_of_images(images)
image_type = get_image_type(images_list[0][0])
images = make_list_of_images(images)
image_type = get_image_type(images[0])
if image_type not in [ImageType.PIL, ImageType.TORCH, ImageType.NUMPY]:
raise ValueError(f"Unsupported input image type {image_type}")
@ -290,65 +319,63 @@ class PixtralImageProcessorFast(BaseImageProcessorFast):
data_format=data_format,
)
if do_convert_rgb:
images_list = [[convert_to_rgb(image) for image in images] for images in images_list]
if image_type == ImageType.PIL:
images_list = [[F.pil_to_tensor(image) for image in images] for images in images_list]
elif image_type == ImageType.NUMPY:
# not using F.to_tensor as it doesn't handle (C, H, W) numpy arrays
images_list = [[torch.from_numpy(image).contiguous() for image in images] for images in images_list]
if device is not None:
images_list = [[image.to(device) for image in images] for images in images_list]
# We assume that all images have the same channel dimension format.
if input_data_format is None:
input_data_format = infer_channel_dimension_format(images_list[0][0])
if input_data_format == ChannelDimension.LAST:
images_list = [[image.permute(2, 0, 1).contiguous() for image in images] for images in images_list]
input_data_format = ChannelDimension.FIRST
if do_rescale and do_normalize:
# fused rescale and normalize
new_mean = torch.tensor(image_mean, device=images_list[0][0].device) * (1.0 / rescale_factor)
new_std = torch.tensor(image_std, device=images_list[0][0].device) * (1.0 / rescale_factor)
new_mean = torch.tensor(image_mean, device=device) * (1.0 / rescale_factor)
new_std = torch.tensor(image_std, device=device) * (1.0 / rescale_factor)
batch_images = []
batch_image_sizes = []
for sample_images in images_list:
images = []
image_sizes = []
for image in sample_images:
if do_resize:
interpolation = (
pil_torch_interpolation_mapping[resample]
if isinstance(resample, (PILImageResampling, int))
else resample
)
image = self.resize(
image=image,
size=size,
patch_size=patch_size,
interpolation=interpolation,
)
for image in images:
if do_convert_rgb:
image = convert_to_rgb(image)
if do_rescale and do_normalize:
# fused rescale and normalize
image = F.normalize(image.to(dtype=torch.float32), new_mean, new_std)
elif do_rescale:
image = image * rescale_factor
elif do_normalize:
image = F.normalize(image, image_mean, image_std)
if image_type == ImageType.PIL:
image = F.pil_to_tensor(image)
elif image_type == ImageType.NUMPY:
# not using F.to_tensor as it doesn't handle (C, H, W) numpy arrays
image = torch.from_numpy(image).contiguous()
images.append(image)
image_sizes.append(get_image_size(image, input_data_format))
batch_images.append(images)
batch_image_sizes.append(image_sizes)
# We assume that all images have the same channel dimension format.
if input_data_format is None:
input_data_format = infer_channel_dimension_format(image)
return BatchMixFeature(
data={"pixel_values": batch_images, "image_sizes": batch_image_sizes},
tensor_type=None,
if input_data_format == ChannelDimension.LAST:
image = image.permute(2, 0, 1).contiguous()
image = image.to(device)
if do_resize:
interpolation = (
pil_torch_interpolation_mapping[resample]
if isinstance(resample, (PILImageResampling, int))
else resample
)
image = self.resize(
image=image,
size=size,
patch_size=patch_size,
interpolation=interpolation,
)
if do_rescale and do_normalize:
# fused rescale and normalize
image = F.normalize(image.to(dtype=torch.float32), new_mean, new_std)
elif do_rescale:
image = image * rescale_factor
elif do_normalize:
image = F.normalize(image, image_mean, image_std)
batch_images.append(image)
batch_image_sizes.append(get_image_size(image, ChannelDimension.FIRST))
pixel_values = self._pad_for_batching(
pixel_values=batch_images,
image_sizes=batch_image_sizes,
)
return BatchFeature(
data={"pixel_values": pixel_values, "image_sizes": batch_image_sizes}, tensor_type=return_tensors
)

View file

@ -14,7 +14,7 @@
# limitations under the License.
"""PyTorch Pixtral model."""
from typing import List, Optional, Tuple, Union
from typing import Optional, Tuple, Union
import torch
import torch.utils.checkpoint
@ -57,7 +57,7 @@ class PixtralRotaryEmbedding(nn.Module):
a corresponding positional embedding, based on its index in the grid.
"""
def __init__(self, config, device):
def __init__(self, config, device=None):
super().__init__()
self.rope_type = "default"
self.dim = config.head_dim
@ -89,7 +89,6 @@ class PixtralRotaryEmbedding(nn.Module):
# Core RoPE block
freqs = self.inv_freq[position_ids]
# position_ids_expanded = position_ids[:, None, :].float()
# Force float32 (see https://github.com/huggingface/transformers/pull/29285)
device_type = x.device.type
device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu"
@ -175,7 +174,7 @@ class PixtralAttention(nn.Module):
self,
hidden_states: torch.Tensor,
attention_mask: Optional[torch.Tensor] = None,
position_embeddings: Optional[torch.Tensor] = None,
position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
output_attentions: Optional[bool] = False,
) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
"""Input shape: Batch x Time x Channel"""
@ -261,8 +260,8 @@ class PixtralAttentionLayer(nn.Module):
self,
hidden_states: torch.Tensor,
attention_mask: torch.Tensor,
position_embeddings: Optional[torch.Tensor] = None,
output_attentions: Optional[bool] = False,
position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
output_attentions: Optional[bool] = None,
) -> Tuple[torch.FloatTensor]:
"""
Args:
@ -310,7 +309,7 @@ class PixtralTransformer(nn.Module):
self,
inputs_embeds,
attention_mask: Optional[torch.Tensor] = None,
position_embeddings: Optional[torch.Tensor] = None,
position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
output_attentions: Optional[bool] = None,
output_hidden_states: Optional[bool] = None,
return_dict: Optional[bool] = None,
@ -375,7 +374,7 @@ class PixtralTransformer(nn.Module):
if not return_dict:
return tuple(v for v in [hidden_states, encoder_states, all_attentions] if v is not None)
return BaseModelOutput(
last_hidden_state=hidden_states, hidden_states=[hidden_states], attentions=all_attentions
last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions
)
@ -399,10 +398,9 @@ PIXTRAL_START_DOCSTRING = r"""
class PixtralPreTrainedModel(PreTrainedModel):
config_class = PixtralVisionConfig
base_model_prefix = "model"
main_input_name = "pixel_values"
supports_gradient_checkpointing = True
_no_split_modules = ["PixtralVisionAttention"]
_skip_keys_device_placement = "past_key_values"
_supports_cache_class = True
_no_split_modules = ["PixtralAttentionLayer"]
def _init_weights(self, module):
std = (
@ -426,6 +424,8 @@ PIXTRAL_INPUTS_DOCSTRING = r"""
pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See [`AutoImageProcessor.__call__`]
for details.
image_sizes (`torch.LongTensor` of shape `(batch_size, 2)`, *optional*):
The sizes of the images in the batch, being (height, width) for each image.
output_attentions (`bool`, *optional*):
Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
tensors for more detail.
@ -470,15 +470,22 @@ class PixtralVisionModel(PixtralPreTrainedModel):
stride=config.patch_size,
bias=False,
)
self.patch_size = config.patch_size
self.ln_pre = PixtralRMSNorm(config.hidden_size, eps=1e-5)
self.transformer = PixtralTransformer(config)
self.patch_positional_embedding = PixtralRotaryEmbedding(config, device=self.device)
self.patch_positional_embedding = PixtralRotaryEmbedding(config)
self.post_init()
def get_input_embeddings(self):
return self.patch_conv
@add_start_docstrings_to_model_forward(PIXTRAL_INPUTS_DOCSTRING)
def forward(
self,
pixel_values: List[torch.Tensor],
output_hidden_states: Optional[bool] = False,
pixel_values: torch.Tensor,
image_sizes: torch.Tensor,
output_hidden_states: Optional[bool] = None,
output_attentions: Optional[bool] = None,
return_dict: Optional[bool] = None,
*args,
@ -490,24 +497,36 @@ class PixtralVisionModel(PixtralPreTrainedModel):
all tokens of all images of shape (N_toks, D)
"""
# pass images through initial convolution independently
if len(pixel_values) > 1:
raise ValueError("Batching/padding not supported yet!")
patch_embeds_list = [self.patch_conv(img.to(self.dtype)) for sample in pixel_values for img in sample]
patch_embeds = self.patch_conv(pixel_values)
patch_embeds_list = [
embed[..., : (size[0] // self.patch_size), : (size[1] // self.patch_size)]
for embed, size in zip(patch_embeds, image_sizes)
]
# flatten to a single sequence
patch_embeds = torch.cat([p.flatten(1).T for p in patch_embeds_list], dim=0).unsqueeze(0)
patch_embeds = self.ln_pre(patch_embeds)
# positional embeddings
position_ids = position_ids_in_meshgrid(
patch_embeds_list, max_width=self.config.image_size // self.config.patch_size
).to(self.device)
position_embedding = self.patch_positional_embedding(patch_embeds, position_ids)
)
position_embeddings = self.patch_positional_embedding(patch_embeds, position_ids)
attention_mask = generate_block_attention_mask(
[p.shape[-2] * p.shape[-1] for p in patch_embeds_list], patch_embeds
)
return self.transformer(patch_embeds, attention_mask, position_embedding)
out = self.transformer(
patch_embeds,
attention_mask=attention_mask,
position_embeddings=position_embeddings,
output_hidden_states=output_hidden_states,
output_attentions=output_attentions,
return_dict=return_dict,
)
return out
__all__ = ["PixtralVisionModel", "PixtralPreTrainedModel"]

View file

@ -22,7 +22,7 @@ from ...feature_extraction_utils import BatchFeature
from ...image_utils import ImageInput, is_valid_image, load_image
from ...processing_utils import ProcessingKwargs, ProcessorMixin, Unpack, _validate_images_text_input_order
from ...tokenization_utils_base import PreTokenizedInput, TextInput
from ...utils import is_torch_device, is_torch_dtype, logging, requires_backends
from ...utils import logging
logger = logging.get_logger(__name__)
@ -50,58 +50,6 @@ def is_image_or_image_url(elem):
return is_url(elem) or is_valid_image(elem)
# Copied from transformers.models.pixtral.image_processing_pixtral.BatchMixFeature
class BatchMixFeature(BatchFeature):
def to(self, *args, **kwargs) -> "BatchMixFeature":
"""
Send all values to device by calling `v.to(*args, **kwargs)` (PyTorch only). This should support casting in
different `dtypes` and sending the `BatchFeature` to a different `device`.
Args:
args (`Tuple`):
Will be passed to the `to(...)` function of the tensors.
kwargs (`Dict`, *optional*):
Will be passed to the `to(...)` function of the tensors.
Returns:
[`BatchFeature`]: The same instance after modification.
"""
def _recursive_to(obj, device, *args, **kwargs):
# Lists can be nested, so keep digging until we hit tensors
if isinstance(obj, list):
return [_recursive_to(o, device, *args, **kwargs) for o in obj]
# We cast only floating point tensors to avoid issues with tokenizers casting `LongTensor` to `FloatTensor`
elif isinstance(obj, torch.Tensor) and torch.is_floating_point(obj):
# cast and send to device
return obj.to(*args, **kwargs)
elif isinstance(obj, torch.Tensor) and device is not None:
# only send to device, don't cast
return obj.to(device=device)
else:
return obj
requires_backends(self, ["torch"])
import torch # noqa
device = kwargs.get("device")
# Check if the args are a device or a dtype
if device is None and len(args) > 0:
# device should be always the first argument
arg = args[0]
if is_torch_dtype(arg):
# The first argument is a dtype
pass
elif isinstance(arg, str) or is_torch_device(arg) or isinstance(arg, int):
device = arg
else:
# it's something else
raise ValueError(f"Attempting to cast a BatchFeature to type {str(arg)}. This is not supported.")
self.data = {k: _recursive_to(v, device, *args, **kwargs) for k, v in self.data.items()}
return self
class PixtralProcessor(ProcessorMixin):
r"""
Constructs a Pixtral processor which wraps a Pixtral image processor and a Pixtral tokenizer into a single processor.
@ -161,7 +109,7 @@ class PixtralProcessor(ProcessorMixin):
audio=None,
videos=None,
**kwargs: Unpack[PixtralProcessorKwargs],
) -> BatchMixFeature:
) -> 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 LlamaTokenizerFast's [`~LlamaTokenizerFast.__call__`] if `text` is not `None` to encode
@ -205,28 +153,16 @@ class PixtralProcessor(ProcessorMixin):
if images is not None:
if is_image_or_image_url(images):
if isinstance(text, str) or isinstance(text, list) and len(text) == 1:
# If there's a single sample, the image must belong to it
images = [[images]]
else:
raise ValueError(
"You have supplied multiple text samples, but `images` is not a nested list. When processing multiple samples, `images` should be a list of lists of images, one list per sample."
)
images = [images]
elif isinstance(images, list) and is_image_or_image_url(images[0]):
if isinstance(text, str) or isinstance(text, list) and len(text) == 1:
# If there's a single sample, all images must belong to it
images = [images]
else:
raise ValueError(
"You have supplied multiple text samples, but `images` is not a nested list. When processing multiple samples, `images` should be a list of lists of images, one list per sample."
)
elif isinstance(images, list) and isinstance(images[0], list) and is_image_or_image_url(images[0][0]):
pass
elif isinstance(images, list) and isinstance(images[0], list) and is_image_or_image_url(images[0][0]):
images = [image for sublist in images for image in sublist]
else:
raise ValueError(
"Invalid input images. Please provide a single image, a list of images, or a list of lists of images."
)
images = [[load_image(im) for im in sample] for sample in images]
images = [load_image(im) if isinstance(im, str) else im for im in images]
image_inputs = self.image_processor(images, patch_size=self.patch_size, **output_kwargs["images_kwargs"])
else:
image_inputs = {}
@ -240,15 +176,13 @@ class PixtralProcessor(ProcessorMixin):
prompt_strings = text
if image_inputs.get("pixel_values") is not None:
# Replace the image token with the expanded image token sequence
images = image_inputs["pixel_values"]
image_sizes = image_inputs.pop("image_sizes")
image_sizes = iter(image_inputs["image_sizes"])
prompt_strings = []
replace_strings = []
for sample_images, sample_image_sizes, sample in zip(images, image_sizes, text):
replace_strings = []
# First calculate the number of tokens needed for each image and put in a placeholder
for image, image_size in zip(sample_images, sample_image_sizes):
height, width = image_size
for sample in text:
while self.image_token in sample:
height, width = next(image_sizes)
num_height_tokens = height // self.patch_size
num_width_tokens = width // self.patch_size
replace_tokens = [
@ -267,7 +201,9 @@ class PixtralProcessor(ProcessorMixin):
prompt_strings.append(sample)
text_inputs = self.tokenizer(prompt_strings, **output_kwargs["text_kwargs"])
return BatchMixFeature(data={**text_inputs, **image_inputs})
return BatchFeature(
data={**text_inputs, **image_inputs}, tensor_type=output_kwargs["common_kwargs"]["return_tensors"]
)
# Copied from transformers.models.clip.processing_clip.CLIPProcessor.batch_decode with CLIP->Llama
def batch_decode(self, *args, **kwargs):

View file

@ -564,9 +564,8 @@ class LlavaForConditionalGenerationIntegrationTest(unittest.TestCase):
self.assertTrue(processor.batch_decode(output, skip_special_tokens=True)[0] == EXPECTED_DECODED_TEXT)
@slow
@require_bitsandbytes
def test_pixtral(self):
model_id = "hf-internal-testing/pixtral-12b"
model_id = "mistral-community/pixtral-12b"
model = LlavaForConditionalGeneration.from_pretrained(model_id)
processor = AutoProcessor.from_pretrained(model_id)
@ -579,33 +578,75 @@ class LlavaForConditionalGenerationIntegrationTest(unittest.TestCase):
PROMPT = "<s>[INST]Describe the images.\n[IMG][IMG][IMG][IMG][/INST]"
# image = Image.open(requests.get(url, stream=True).raw)
inputs = processor(text=PROMPT, images=IMG_URLS, return_tensors="pt").to("cuda")
inputs = processor(text=PROMPT, images=IMG_URLS, return_tensors="pt").to(model.device)
generate_ids = model.generate(**inputs, max_new_tokens=500)
ouptut = processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
print(ouptut)
# fmt: off
EXPECTED_GENERATION = """
Describe the images.
Sure, let's break down each image description:
Certainly! Here are the descriptions of the images:
1. **Image 1:**
- **Description:** A black dog with a glossy coat is sitting on a wooden floor. The dog has a focused expression and is looking directly at the camera.
- **Details:** The wooden floor has a rustic appearance with visible wood grain patterns. The dog's eyes are a striking color, possibly brown or amber, which contrasts with its black fur.
1. **Image 1**: This image features a black dog with a glossy coat sitting on a wooden surface. The dog has a calm and attentive expression, looking directly at the camera. The wooden background has a rustic appearance with visible grain and texture.
2. **Image 2:**
- **Description:** A scenic view of a mountainous landscape with a winding road cutting through it. The road is surrounded by lush green vegetation and leads to a distant valley.
- **Details:** The mountains are rugged with steep slopes, and the sky is clear, indicating good weather. The winding road adds a sense of depth and perspective to the image.
2. **Image 2**: This image captures a breathtaking view of a mountainous landscape. The mountains are rugged and covered with patches of green vegetation. The sky above is clear, and the scene conveys a sense of tranquility and natural beauty.
3. **Image 3:**
- **Description:** A beach scene with waves crashing against the shore. There are several people in the water and on the beach, enjoying the waves and the sunset.
- **Details:** The waves are powerful, creating a dynamic and lively atmosphere. The sky is painted with hues of orange and pink from the setting sun, adding a warm glow to the scene.
3. **Image 3**: This image shows a beach scene during sunset. The waves are gently rolling onto the shore, and several people can be seen in the water, possibly surfing or swimming. The sky is painted with warm hues of orange and yellow, creating a serene and picturesque atmosphere.
4. **Image 4:**
- **Description:** A garden path leading to a large tree with a bench underneath it. The path is bordered by well-maintained grass and flowers.
- **Details:** The path is made of small stones or gravel, and the tree provides a shaded area with the bench invitingly placed beneath it. The surrounding area is lush and green, suggesting a well-kept garden.
4. **Image 4**: This image depicts a narrow, winding path that cuts through a lush, green landscape. On either side of the path, there is dense grass and various trees, including a prominent tree with white blossoms. The sky is clear and blue, adding to the peaceful and inviting ambiance of the scene.
Each image captures a different scene, from a close-up of a dog to expansive natural landscapes, showcasing various elements of nature and human interaction with it.
These descriptions provide a detailed overview of the content and atmosphere of each image.
"""
# fmt: on
# check that both inputs are handled correctly and generate the same output
self.assertListEqual(ouptut, EXPECTED_GENERATION)
self.assertEqual(ouptut, EXPECTED_GENERATION)
@slow
@require_bitsandbytes
def test_pixtral_4bit(self):
model_id = "mistral-community/pixtral-12b"
model = LlavaForConditionalGeneration.from_pretrained(model_id, load_in_4bit=True)
processor = AutoProcessor.from_pretrained(model_id)
IMG_URLS = [
Image.open(requests.get("https://picsum.photos/id/237/400/300", stream=True).raw),
Image.open(requests.get("https://picsum.photos/id/231/200/300", stream=True).raw),
]
PROMPT = "<s>[INST][IMG][IMG]Describe the images.[/INST]"
inputs = processor(text=PROMPT, images=IMG_URLS, return_tensors="pt").to(torch_device, torch.float16)
generate_ids = model.generate(**inputs, max_new_tokens=50)
output = processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
EXPECTED_GENERATION = "Describe the images.The image showcases a dog, which is prominently positioned in the center, taking up a significant portion of the frame. The dog is situated against a backdrop of a wooden surface, which spans the entire image. The dog appears to be a black Labrador" # fmt: skip
self.assertEqual(output, EXPECTED_GENERATION)
@slow
@require_bitsandbytes
def test_pixtral_batched(self):
model_id = "mistral-community/pixtral-12b"
model = LlavaForConditionalGeneration.from_pretrained(model_id, load_in_4bit=True)
processor = AutoProcessor.from_pretrained(model_id)
processor.tokenizer.pad_token_id = processor.tokenizer.eos_token_id
IMG_URLS = [
Image.open(requests.get("https://picsum.photos/id/237/400/300", stream=True).raw),
Image.open(requests.get("https://picsum.photos/id/17/150/500", stream=True).raw),
]
PROMPT = [
"<s>[INST][IMG]What breed is the dog?[/INST]",
"<s>[INST][IMG]What is shown in this image?[/INST]",
]
inputs = processor(text=PROMPT, images=IMG_URLS, padding=True, return_tensors="pt").to(
torch_device, torch.float16
)
generate_ids = model.generate(**inputs, max_new_tokens=50)
output = processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)
EXPECTED_GENERATION = [
'What breed is the dog?The dog in the image is a black Labrador Retriever.',
'What is shown in this image?The image depicts a narrow, winding dirt path surrounded by lush greenery. The path is flanked by grass and shrubs on both sides. On the left side, there are tall trees and dense foliage, while on the right side, there'
] # fmt: skip
self.assertEqual(output, EXPECTED_GENERATION)

View file

@ -13,7 +13,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import random
import time
import unittest
@ -92,49 +91,47 @@ class PixtralImageProcessingTester:
"do_convert_rgb": self.do_convert_rgb,
}
def expected_output_image_shape(self, image):
if isinstance(image, Image.Image):
width, height = image.size
elif isinstance(image, np.ndarray):
height, width = image.shape[:2]
elif isinstance(image, torch.Tensor):
height, width = image.shape[-2:]
def expected_output_image_shape(self, images):
if not isinstance(images, (list, tuple)):
images = [images]
max_height = max_width = self.size.get("longest_edge")
batch_size = len(images)
return_height, return_width = 0, 0
for image in images:
if isinstance(image, Image.Image):
width, height = image.size
elif isinstance(image, np.ndarray):
height, width = image.shape[:2]
elif isinstance(image, torch.Tensor):
height, width = image.shape[-2:]
ratio = max(height / max_height, width / max_width)
if ratio > 1:
height = int(np.ceil(height / ratio))
width = int(np.ceil(width / ratio))
max_height = max_width = self.size.get("longest_edge")
patch_height, patch_width = self.patch_size["height"], self.patch_size["width"]
num_height_tokens = (height - 1) // patch_height + 1
num_width_tokens = (width - 1) // patch_width + 1
ratio = max(height / max_height, width / max_width)
if ratio > 1:
height = int(np.ceil(height / ratio))
width = int(np.ceil(width / ratio))
height = num_height_tokens * patch_height
width = num_width_tokens * patch_width
patch_height, patch_width = self.patch_size["height"], self.patch_size["width"]
num_height_tokens = (height - 1) // patch_height + 1
num_width_tokens = (width - 1) // patch_width + 1
return self.num_channels, height, width
return_height = max(num_height_tokens * patch_height, return_height)
return_width = max(num_width_tokens * patch_width, return_width)
return batch_size, self.num_channels, return_height, return_width
def prepare_image_inputs(self, equal_resolution=False, numpify=False, torchify=False):
# Use prepare_image_inputs to make a list of list of single images
images_list = []
for _ in range(self.batch_size):
images = []
for _ in range(random.randint(1, self.max_num_images_per_sample)):
img = prepare_image_inputs(
batch_size=1,
num_channels=self.num_channels,
min_resolution=self.min_resolution,
max_resolution=self.max_resolution,
equal_resolution=equal_resolution,
numpify=numpify,
torchify=torchify,
)[0]
images.append(img)
images_list.append(images)
return images_list
images = prepare_image_inputs(
batch_size=self.batch_size,
num_channels=self.num_channels,
min_resolution=self.min_resolution,
max_resolution=self.max_resolution,
equal_resolution=equal_resolution,
numpify=numpify,
torchify=torchify,
)
return images
@require_torch
@ -173,23 +170,18 @@ class PixtralImageProcessingTest(ImageProcessingTestMixin, unittest.TestCase):
image_processing = image_processing_class(**self.image_processor_dict)
# create random PIL images
image_inputs_list = self.image_processor_tester.prepare_image_inputs()
for image_inputs in image_inputs_list:
for image in image_inputs:
self.assertIsInstance(image, Image.Image)
for image in image_inputs_list:
self.assertIsInstance(image, Image.Image)
# Test not batched input
encoded_images = image_processing(image_inputs_list[0][0], return_tensors="pt").pixel_values
expected_output_image_shape = self.image_processor_tester.expected_output_image_shape(
image_inputs_list[0][0]
)
self.assertEqual(tuple(encoded_images[0][0].shape), expected_output_image_shape)
encoded_images = image_processing(image_inputs_list[0], return_tensors="pt").pixel_values
expected_output_image_shape = self.image_processor_tester.expected_output_image_shape(image_inputs_list[0])
self.assertEqual(tuple(encoded_images.shape), expected_output_image_shape)
# Test batched
batch_encoded_images = image_processing(image_inputs_list, return_tensors="pt").pixel_values
for encoded_images, images in zip(batch_encoded_images, image_inputs_list):
for encoded_image, image in zip(encoded_images, images):
expected_output_image_shape = self.image_processor_tester.expected_output_image_shape(image)
self.assertEqual(tuple(encoded_image.shape), expected_output_image_shape)
encoded_images = image_processing(image_inputs_list, return_tensors="pt").pixel_values
expected_output_image_shape = self.image_processor_tester.expected_output_image_shape(image_inputs_list)
self.assertEqual(tuple(encoded_images.shape), expected_output_image_shape)
def test_call_numpy(self):
for image_processing_class in self.image_processor_list:
@ -197,23 +189,18 @@ class PixtralImageProcessingTest(ImageProcessingTestMixin, unittest.TestCase):
image_processing = image_processing_class(**self.image_processor_dict)
# create random numpy tensors
image_inputs_list = self.image_processor_tester.prepare_image_inputs(numpify=True)
for image_inputs in image_inputs_list:
for image in image_inputs:
self.assertIsInstance(image, np.ndarray)
for image in image_inputs_list:
self.assertIsInstance(image, np.ndarray)
# Test not batched input
encoded_images = image_processing(image_inputs_list[0][0], return_tensors="pt").pixel_values
expected_output_image_shape = self.image_processor_tester.expected_output_image_shape(
image_inputs_list[0][0]
)
self.assertEqual(tuple(encoded_images[0][0].shape), expected_output_image_shape)
encoded_images = image_processing(image_inputs_list[0], return_tensors="pt").pixel_values
expected_output_image_shape = self.image_processor_tester.expected_output_image_shape(image_inputs_list[0])
self.assertEqual(tuple(encoded_images.shape), expected_output_image_shape)
# Test batched
batch_encoded_images = image_processing(image_inputs_list, return_tensors="pt").pixel_values
for encoded_images, images in zip(batch_encoded_images, image_inputs_list):
for encoded_image, image in zip(encoded_images, images):
expected_output_image_shape = self.image_processor_tester.expected_output_image_shape(image)
self.assertEqual(tuple(encoded_image.shape), expected_output_image_shape)
expected_output_image_shape = self.image_processor_tester.expected_output_image_shape(image_inputs_list)
self.assertEqual(tuple(batch_encoded_images.shape), expected_output_image_shape)
def test_call_pytorch(self):
for image_processing_class in self.image_processor_list:
@ -221,23 +208,18 @@ class PixtralImageProcessingTest(ImageProcessingTestMixin, unittest.TestCase):
image_processing = image_processing_class(**self.image_processor_dict)
# create random PyTorch tensors
image_inputs_list = self.image_processor_tester.prepare_image_inputs(torchify=True)
for image_inputs in image_inputs_list:
for image in image_inputs:
self.assertIsInstance(image, torch.Tensor)
for image in image_inputs_list:
self.assertIsInstance(image, torch.Tensor)
# Test not batched input
encoded_images = image_processing(image_inputs_list[0][0], return_tensors="pt").pixel_values
expected_output_image_shape = self.image_processor_tester.expected_output_image_shape(
image_inputs_list[0][0]
)
self.assertEqual(tuple(encoded_images[0][0].shape), expected_output_image_shape)
encoded_images = image_processing(image_inputs_list[0], return_tensors="pt").pixel_values
expected_output_image_shape = self.image_processor_tester.expected_output_image_shape(image_inputs_list[0])
self.assertEqual(tuple(encoded_images.shape), expected_output_image_shape)
# Test batched
batch_encoded_images = image_processing(image_inputs_list, return_tensors="pt").pixel_values
for encoded_images, images in zip(batch_encoded_images, image_inputs_list):
for encoded_image, image in zip(encoded_images, images):
expected_output_image_shape = self.image_processor_tester.expected_output_image_shape(image)
self.assertEqual(tuple(encoded_image.shape), expected_output_image_shape)
expected_output_image_shape = self.image_processor_tester.expected_output_image_shape(image_inputs_list)
self.assertEqual(tuple(batch_encoded_images.shape), expected_output_image_shape)
@require_vision
@require_torch

View file

@ -74,15 +74,17 @@ class PixtralVisionModelTester:
self.initializer_range = initializer_range
self.scope = scope
# in ViT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
num_patches = (image_size // patch_size) ** 2
self.seq_length = num_patches + 1
# in Pixtral, the seq length equals the number of patches * batch_size because the patches are flattened
self.seq_length = (image_size // patch_size) ** 2 * batch_size
def prepare_config_and_inputs(self):
pixel_values = floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size])
image_sizes = torch.tensor(
[[self.image_size, self.image_size]] * self.batch_size, dtype=torch.long, device=torch_device
)
config = self.get_config()
return config, pixel_values
return config, pixel_values, image_sizes
def get_config(self):
return PixtralVisionConfig(
@ -127,8 +129,8 @@ class PixtralVisionModelTester:
def prepare_config_and_inputs_for_common(self):
config_and_inputs = self.prepare_config_and_inputs()
config, pixel_values = config_and_inputs
inputs_dict = {"pixel_values": pixel_values}
config, pixel_values, image_sizes = config_and_inputs
inputs_dict = {"pixel_values": pixel_values, "image_sizes": image_sizes}
return config, inputs_dict
@ -142,113 +144,17 @@ class PixtralVisionModelModelTest(ModelTesterMixin, unittest.TestCase):
test_pruning = False
test_head_masking = False
test_torchscript = False
test_resize_embeddings = False
def setUp(self):
self.model_tester = PixtralVisionModelTester(self)
self.config_tester = ConfigTester(self, config_class=PixtralVisionConfig, has_text_modality=False)
@unittest.skip("model does not support input embeds")
def test_inputs_embeds(self):
pass
@unittest.skip("model does not support input embeds")
def test_inputs_embeds_matches_input_ids(self):
pass
@unittest.skip(
reason="This architecure seem to not compute gradients properly when using GC, check: https://github.com/huggingface/transformers/pull/27124"
)
def test_training_gradient_checkpointing(self):
pass
@unittest.skip(
reason="This architecure seem to not compute gradients properly when using GC, check: https://github.com/huggingface/transformers/pull/27124"
)
def test_training_gradient_checkpointing_use_reentrant(self):
pass
@unittest.skip(
reason="This architecure seem to not compute gradients properly when using GC, check: https://github.com/huggingface/transformers/pull/27124"
)
def test_training_gradient_checkpointing_use_reentrant_false(self):
pass
@unittest.skip(reason="Compile not yet supported because in Pixtral models")
def test_sdpa_can_compile_dynamic(self):
pass
@unittest.skip(reason="Compile not yet supported because in Pixtral models")
def test_sdpa_can_dispatch_on_flash(self):
pass
@unittest.skip(reason="Not supported yet")
def test_attention_outputs(self):
pass
@unittest.skip(reason="Not supported yet")
def test_cpu_offload(self):
pass
@unittest.skip(reason="Not supported yet")
def test_batching_equivalence(self):
pass
@unittest.skip(reason="Not supported yet")
def test_disk_offload_bin(self):
pass
@unittest.skip(reason="Not supported yet")
def test_retain_grad_hidden_states_attentions(self):
pass
@unittest.skip(reason="Not supported yet")
def test_multi_gpu_data_parallel_forward(self):
pass
@unittest.skip(reason="Not supported yet")
def test_model_parallelism(self):
pass
@unittest.skip(reason="Not supported yet")
def test_model_outputs_equivalence(self):
pass
@unittest.skip(reason="Not supported yet")
def test_save_load(self):
pass
@unittest.skip(reason="Not supported yet")
def test_model_get_set_embeddings(self):
pass
config, _ = self.model_tester.prepare_config_and_inputs_for_common()
@unittest.skip(reason="Not supported yet")
def test_resize_tokens_embeddings(self):
pass
@unittest.skip(reason="Not supported yet")
def test_model_main_input_name(self):
pass
@unittest.skip(reason="Not supported yet")
def test_initialization(self):
pass
@unittest.skip(reason="Not supported yet")
def test_hidden_states_output(self):
pass
@unittest.skip(reason="Not supported yet")
def test_gradient_checkpointing_backward_compatibility(self):
pass
@unittest.skip(reason="Not supported yet")
def test_feed_forward_chunking(self):
pass
@unittest.skip(reason="Not supported yet")
def test_disk_offload_safetensors(self):
pass
@unittest.skip(reason="Not supported yet")
def test_determinism(self):
pass
for model_class in self.all_model_classes:
model = model_class(config)
self.assertIsInstance(model.get_input_embeddings(), (torch.nn.Module))
x = model.get_output_embeddings()
self.assertTrue(x is None or isinstance(x, torch.nn.Linear))

View file

@ -14,7 +14,6 @@
import shutil
import tempfile
import unittest
from typing import Optional
import requests
import torch
@ -28,7 +27,7 @@ from ...test_processing_common import ProcessorTesterMixin
if is_vision_available():
from PIL import Image
from transformers import AutoTokenizer, PixtralImageProcessor, PixtralProcessor
from transformers import PixtralProcessor
@require_vision
@ -46,20 +45,15 @@ class PixtralProcessorTest(ProcessorTesterMixin, unittest.TestCase):
def setUp(self):
self.tmpdirname = tempfile.mkdtemp()
# FIXME - just load the processor directly from the checkpoint
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/pixtral-12b")
image_processor = PixtralImageProcessor()
processor = PixtralProcessor(tokenizer=tokenizer, image_processor=image_processor)
processor = PixtralProcessor.from_pretrained("mistral-community/pixtral-12b")
processor.save_pretrained(self.tmpdirname)
def tearDown(self):
shutil.rmtree(self.tmpdirname)
@unittest.skip("No chat template was set for this model (yet)")
def test_chat_template(self):
processor = self.processor_class.from_pretrained(self.tmpdirname)
expected_prompt = "USER: [IMG]\nWhat is shown in this image? ASSISTANT:"
expected_prompt = "<s>[INST][IMG]What is shown in this image?[/INST]"
messages = [
{
@ -73,13 +67,12 @@ class PixtralProcessorTest(ProcessorTesterMixin, unittest.TestCase):
formatted_prompt = processor.apply_chat_template(messages, add_generation_prompt=True)
self.assertEqual(expected_prompt, formatted_prompt)
@unittest.skip("No chat template was set for this model (yet)")
def test_image_token_filling(self):
processor = self.processor_class.from_pretrained(self.tmpdirname)
# Important to check with non square image
image = torch.randint(0, 2, (3, 500, 316))
expected_image_tokens = 1526
image_token_index = 32000
expected_image_tokens = 640
image_token_index = 10
messages = [
{
@ -111,11 +104,8 @@ class PixtralProcessorTest(ProcessorTesterMixin, unittest.TestCase):
self.assertIn("input_ids", inputs_image)
self.assertTrue(len(inputs_image["input_ids"]) == 1)
self.assertIsInstance(inputs_image["input_ids"], torch.Tensor)
self.assertIsInstance(inputs_image["pixel_values"], list)
self.assertTrue(len(inputs_image["pixel_values"]) == 1)
self.assertIsInstance(inputs_image["pixel_values"][0], list)
self.assertTrue(len(inputs_image["pixel_values"][0]) == 1)
self.assertIsInstance(inputs_image["pixel_values"][0][0], torch.Tensor)
self.assertIsInstance(inputs_image["pixel_values"], torch.Tensor)
self.assertTrue(inputs_image["pixel_values"].shape == torch.Size([1, 3, 32, 32]))
# fmt: off
input_ids = inputs_image["input_ids"]
@ -131,11 +121,8 @@ class PixtralProcessorTest(ProcessorTesterMixin, unittest.TestCase):
self.assertIn("input_ids", inputs_url)
self.assertTrue(len(inputs_url["input_ids"]) == 1)
self.assertIsInstance(inputs_url["input_ids"], torch.Tensor)
self.assertIsInstance(inputs_url["pixel_values"], list)
self.assertTrue(len(inputs_url["pixel_values"]) == 1)
self.assertIsInstance(inputs_url["pixel_values"][0], list)
self.assertTrue(len(inputs_url["pixel_values"][0]) == 1)
self.assertIsInstance(inputs_url["pixel_values"][0][0], torch.Tensor)
self.assertIsInstance(inputs_image["pixel_values"], torch.Tensor)
self.assertTrue(inputs_image["pixel_values"].shape == torch.Size([1, 3, 32, 32]))
# fmt: off
input_ids = inputs_url["input_ids"]
@ -146,6 +133,28 @@ class PixtralProcessorTest(ProcessorTesterMixin, unittest.TestCase):
)
# fmt: on
# Test passing inputs as a single list
inputs_image = processor(text=prompt_string, images=[self.image_0], return_tensors="pt")
self.assertTrue(inputs_image["pixel_values"].shape == torch.Size([1, 3, 32, 32]))
# fmt: off
self.assertEqual(
inputs_image["input_ids"][0].tolist(),
[21510, 1058, 1032, 10, 10, 12, 10, 10, 13, 1010, 7493, 1681, 1278, 4701, 1307, 1278, 3937, 1063, 1349, 4290, 16002, 41150, 1058]
)
# fmt: on
# Test as nested single list
inputs_image = processor(text=prompt_string, images=[[self.image_0]], return_tensors="pt")
self.assertTrue(inputs_image["pixel_values"].shape == torch.Size([1, 3, 32, 32]))
# fmt: off
self.assertEqual(
inputs_image["input_ids"][0].tolist(),
[21510, 1058, 1032, 10, 10, 12, 10, 10, 13, 1010, 7493, 1681, 1278, 4701, 1307, 1278, 3937, 1063, 1349, 4290, 16002, 41150, 1058]
)
# fmt: on
def test_processor_with_multiple_images_single_list(self):
processor = self.processor_class.from_pretrained(self.tmpdirname)
prompt_string = "USER: [IMG][IMG]\nWhat's the difference between these two images? ASSISTANT:"
@ -159,11 +168,8 @@ class PixtralProcessorTest(ProcessorTesterMixin, unittest.TestCase):
self.assertIn("input_ids", inputs_image)
self.assertTrue(len(inputs_image["input_ids"]) == 1)
self.assertIsInstance(inputs_image["input_ids"], torch.Tensor)
self.assertIsInstance(inputs_image["pixel_values"], list)
self.assertTrue(len(inputs_image["pixel_values"]) == 1)
self.assertIsInstance(inputs_image["pixel_values"][0], list)
self.assertTrue(len(inputs_image["pixel_values"][0]) == 2)
self.assertIsInstance(inputs_image["pixel_values"][0][0], torch.Tensor)
self.assertIsInstance(inputs_image["pixel_values"], torch.Tensor)
self.assertTrue(inputs_image["pixel_values"].shape == torch.Size([2, 3, 32, 32]))
# fmt: off
input_ids = inputs_image["input_ids"]
@ -179,11 +185,9 @@ class PixtralProcessorTest(ProcessorTesterMixin, unittest.TestCase):
self.assertIn("input_ids", inputs_url)
self.assertTrue(len(inputs_url["input_ids"]) == 1)
self.assertIsInstance(inputs_url["input_ids"], torch.Tensor)
self.assertIsInstance(inputs_url["pixel_values"], list)
self.assertTrue(len(inputs_url["pixel_values"]) == 1)
self.assertIsInstance(inputs_url["pixel_values"][0], list)
self.assertTrue(len(inputs_url["pixel_values"][0]) == 2)
self.assertIsInstance(inputs_url["pixel_values"][0][0], torch.Tensor)
self.assertIsInstance(inputs_image["pixel_values"], torch.Tensor)
self.assertTrue(inputs_image["pixel_values"].shape == torch.Size([2, 3, 32, 32]))
# fmt: off
input_ids = inputs_url["input_ids"]
self.assertEqual(
@ -193,6 +197,17 @@ class PixtralProcessorTest(ProcessorTesterMixin, unittest.TestCase):
)
# fmt: on
# Test passing in as a nested list
inputs_url = processor(text=prompt_string, images=[[self.image_0, self.image_1]], return_tensors="pt")
self.assertTrue(inputs_image["pixel_values"].shape == torch.Size([2, 3, 32, 32]))
# fmt: off
self.assertEqual(
inputs_url["input_ids"][0].tolist(),
[21510, 1058, 1032, 10, 10, 12, 10, 10, 13, 10, 10, 12, 10, 10, 13, 1010, 7493, 1681, 1278, 6592, 2396, 2576, 2295, 8061, 1063, 1349, 4290, 16002, 41150, 1058]
)
# fmt: on
def test_processor_with_multiple_images_multiple_lists(self):
processor = self.processor_class.from_pretrained(self.tmpdirname)
prompt_string = [
@ -211,11 +226,8 @@ class PixtralProcessorTest(ProcessorTesterMixin, unittest.TestCase):
self.assertIn("input_ids", inputs_image)
self.assertTrue(len(inputs_image["input_ids"]) == 2)
self.assertIsInstance(inputs_image["input_ids"], torch.Tensor)
self.assertIsInstance(inputs_image["pixel_values"], list)
self.assertTrue(len(inputs_image["pixel_values"]) == 2)
self.assertIsInstance(inputs_image["pixel_values"][0], list)
self.assertTrue(len(inputs_image["pixel_values"][0]) == 2)
self.assertIsInstance(inputs_image["pixel_values"][0][0], torch.Tensor)
self.assertIsInstance(inputs_image["pixel_values"], torch.Tensor)
self.assertTrue(inputs_image["pixel_values"].shape == torch.Size([3, 3, 32, 32]))
# fmt: off
input_ids = inputs_image["input_ids"]
@ -231,11 +243,8 @@ class PixtralProcessorTest(ProcessorTesterMixin, unittest.TestCase):
self.assertIn("input_ids", inputs_url)
self.assertTrue(len(inputs_url["input_ids"]) == 2)
self.assertIsInstance(inputs_url["input_ids"], torch.Tensor)
self.assertIsInstance(inputs_url["pixel_values"], list)
self.assertTrue(len(inputs_url["pixel_values"]) == 2)
self.assertIsInstance(inputs_url["pixel_values"][0], list)
self.assertTrue(len(inputs_url["pixel_values"][0]) == 2)
self.assertIsInstance(inputs_url["pixel_values"][0][0], torch.Tensor)
self.assertIsInstance(inputs_image["pixel_values"], torch.Tensor)
self.assertTrue(inputs_image["pixel_values"].shape == torch.Size([3, 3, 32, 32]))
# fmt: off
input_ids = inputs_url["input_ids"]
@ -246,6 +255,19 @@ class PixtralProcessorTest(ProcessorTesterMixin, unittest.TestCase):
)
# fmt: on
# Test passing as a single flat list
inputs_image = processor(
text=prompt_string, images=[self.image_0, self.image_1, self.image_2], return_tensors="pt", padding=True
)
self.assertTrue(inputs_image["pixel_values"].shape == torch.Size([3, 3, 32, 32]))
# fmt: off
self.assertEqual(
inputs_image["input_ids"][0].tolist(),
[21510, 1058, 1032, 10, 10, 12, 10, 10, 13, 10, 10, 12, 10, 10, 13, 1010, 7493, 1681, 1278, 6592, 2396, 2576, 2295, 8061, 1063, 1349, 4290, 16002, 41150, 1058]
)
# fmt: on
def test_processor_returns_full_length_batches(self):
# to avoid https://github.com/huggingface/transformers/issues/34204
processor = self.processor_class.from_pretrained(self.tmpdirname)
@ -264,13 +286,3 @@ class PixtralProcessorTest(ProcessorTesterMixin, unittest.TestCase):
self.assertIn("input_ids", inputs_image)
self.assertTrue(len(inputs_image["input_ids"]) == 5)
self.assertTrue(len(inputs_image["pixel_values"]) == 5)
# Override as PixtralProcessor needs nested images to work properly with batched inputs
@require_vision
def prepare_image_inputs(self, batch_size: Optional[int] = None):
"""This function prepares a list of PIL images for testing"""
if batch_size is None:
return super().prepare_image_inputs()
if batch_size < 1:
raise ValueError("batch_size must be greater than 0")
return [[super().prepare_image_inputs()]] * batch_size

View file

@ -2991,6 +2991,10 @@ class ModelTesterMixin:
model.to(torch_device)
model.eval()
model_forward_args = inspect.signature(model.forward).parameters
if "inputs_embeds" not in model_forward_args:
self.skipTest(reason="This model doesn't use `inputs_embeds`")
inputs = copy.deepcopy(self._prepare_for_class(inputs_dict, model_class))
if not self.is_encoder_decoder: