|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
from typing import Optional, Union |
|
|
|
|
|
import numpy as np |
|
|
|
|
|
from ...image_processing_utils import BatchFeature |
|
|
from ...image_utils import ImageInput |
|
|
from ...processing_utils import MultiModalData, ProcessingKwargs, ProcessorMixin, Unpack |
|
|
from ...tokenization_utils import PreTokenizedInput, TextInput |
|
|
from ...utils import TensorType |
|
|
from ..auto import AutoTokenizer |
|
|
|
|
|
|
|
|
class AriaProcessorKwargs(ProcessingKwargs, total=False): |
|
|
_defaults = { |
|
|
"text_kwargs": { |
|
|
"padding": False, |
|
|
"return_mm_token_type_ids": False, |
|
|
}, |
|
|
"images_kwargs": { |
|
|
"max_image_size": 980, |
|
|
"split_image": False, |
|
|
}, |
|
|
"return_tensors": TensorType.PYTORCH, |
|
|
} |
|
|
|
|
|
|
|
|
class AriaProcessor(ProcessorMixin): |
|
|
""" |
|
|
AriaProcessor is a processor for the Aria model which wraps the Aria image preprocessor and the LLama slow tokenizer. |
|
|
|
|
|
Args: |
|
|
image_processor (`AriaImageProcessor`, *optional*): |
|
|
The AriaImageProcessor to use for image preprocessing. |
|
|
tokenizer (`PreTrainedTokenizerBase`, *optional*): |
|
|
An instance of [`PreTrainedTokenizerBase`]. This should correspond with the model's text model. 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. |
|
|
size_conversion (`Dict`, *optional*): |
|
|
A dictionary indicating size conversions for images. |
|
|
""" |
|
|
|
|
|
attributes = ["image_processor", "tokenizer"] |
|
|
image_processor_class = "AriaImageProcessor" |
|
|
tokenizer_class = "AutoTokenizer" |
|
|
|
|
|
def __init__( |
|
|
self, |
|
|
image_processor=None, |
|
|
tokenizer: Union[AutoTokenizer, str] = None, |
|
|
chat_template: Optional[str] = None, |
|
|
size_conversion: Optional[dict[Union[float, int], int]] = None, |
|
|
): |
|
|
if size_conversion is None: |
|
|
size_conversion = {490: 128, 980: 256} |
|
|
self.size_conversion = {int(k): v for k, v in size_conversion.items()} |
|
|
|
|
|
self.image_token = tokenizer.image_token |
|
|
self.image_token_id = tokenizer.image_token_id |
|
|
if tokenizer is not None and tokenizer.pad_token is None: |
|
|
tokenizer.pad_token = tokenizer.unk_token |
|
|
|
|
|
super().__init__(image_processor, tokenizer, chat_template=chat_template) |
|
|
|
|
|
def __call__( |
|
|
self, |
|
|
text: Union[TextInput, PreTokenizedInput, list[TextInput], list[PreTokenizedInput]], |
|
|
images: Optional[ImageInput] = None, |
|
|
audio=None, |
|
|
videos=None, |
|
|
**kwargs: Unpack[AriaProcessorKwargs], |
|
|
) -> BatchFeature: |
|
|
""" |
|
|
Main method to prepare for the model one or several sequences(s) and image(s). |
|
|
|
|
|
Args: |
|
|
text (`TextInput`, `PreTokenizedInput`, `list[TextInput]`, `list[PreTokenizedInput]`): |
|
|
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). |
|
|
images (`ImageInput`): |
|
|
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. |
|
|
|
|
|
|
|
|
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_mask** -- Pixel mask to be fed to a model. Returned when `images` is not `None`. |
|
|
""" |
|
|
output_kwargs = self._merge_kwargs( |
|
|
AriaProcessorKwargs, |
|
|
tokenizer_init_kwargs=self.tokenizer.init_kwargs, |
|
|
**kwargs, |
|
|
) |
|
|
|
|
|
if isinstance(text, str): |
|
|
text = [text] |
|
|
elif not isinstance(text, list) and not isinstance(text[0], str): |
|
|
raise TypeError("Invalid input text. Please provide a string, or a list of strings") |
|
|
|
|
|
if images is not None: |
|
|
image_inputs = self.image_processor(images, **output_kwargs["images_kwargs"]) |
|
|
|
|
|
tokens_per_image = self.size_conversion[image_inputs.pixel_values.shape[2]] |
|
|
prompt_strings = [] |
|
|
num_crops = image_inputs.pop("num_crops") * tokens_per_image |
|
|
for sample in text: |
|
|
sample = sample.replace(self.tokenizer.image_token, self.tokenizer.image_token * num_crops) |
|
|
prompt_strings.append(sample) |
|
|
|
|
|
else: |
|
|
image_inputs = {} |
|
|
prompt_strings = text |
|
|
|
|
|
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", False) |
|
|
text_inputs = self.tokenizer(prompt_strings, **output_kwargs["text_kwargs"], return_tensors=None) |
|
|
self._check_special_mm_tokens(prompt_strings, text_inputs, modalities=["image"]) |
|
|
|
|
|
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}, tensor_type=return_tensors) |
|
|
|
|
|
def _get_num_multimodal_tokens(self, image_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. |
|
|
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 = AriaProcessorKwargs._defaults.get("images_kwargs", {}) |
|
|
images_kwargs.update(kwargs) |
|
|
|
|
|
max_size = images_kwargs.get("max_image_size", None) or self.image_processor.max_image_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 = [self.size_conversion[max_size] * num_patches 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) |
|
|
|
|
|
@property |
|
|
def model_input_names(self): |
|
|
tokenizer_input_names = self.tokenizer.model_input_names |
|
|
image_processor_input_names = self.image_processor.model_input_names |
|
|
|
|
|
|
|
|
|
|
|
image_processor_input_names = [name for name in image_processor_input_names if name != "num_crops"] |
|
|
return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names)) |
|
|
|
|
|
|
|
|
__all__ = ["AriaProcessor"] |
|
|
|
