Aluode/PerceptionLabPortable

1# coding=utf-82# Copyright 2023 The HuggingFace Inc. team.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8#     http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15"""16Processor class for Pix2Struct.17"""18 19from typing import Optional, Union20 21from ...feature_extraction_utils import BatchFeature22from ...processing_utils import ImagesKwargs, ProcessingKwargs, ProcessorMixin, Unpack23from ...tokenization_utils_base import BatchEncoding, PreTokenizedInput, TextInput24from ...utils import logging25 26 27class Pix2StructImagesKwargs(ImagesKwargs, total=False):28    max_patches: Optional[int]29    header_text: Optional[Union[TextInput, PreTokenizedInput, list[TextInput], list[PreTokenizedInput]]]30 31 32class Pix2StructProcessorKwargs(ProcessingKwargs, total=False):33    images_kwargs: Pix2StructImagesKwargs34    _defaults = {35        "text_kwargs": {36            "add_special_tokens": True,37            "padding": False,38            "stride": 0,39            "return_overflowing_tokens": False,40            "return_special_tokens_mask": False,41            "return_offsets_mapping": False,42            "return_token_type_ids": False,43            "return_length": False,44            "verbose": True,45        },46        "images_kwargs": {47            "max_patches": 2048,48        },49    }50 51 52logger = logging.get_logger(__name__)53 54 55class Pix2StructProcessor(ProcessorMixin):56    r"""57    Constructs a PIX2STRUCT processor which wraps a BERT tokenizer and PIX2STRUCT image processor into a single58    processor.59 60    [`Pix2StructProcessor`] offers all the functionalities of [`Pix2StructImageProcessor`] and [`T5TokenizerFast`]. See61    the docstring of [`~Pix2StructProcessor.__call__`] and [`~Pix2StructProcessor.decode`] for more information.62 63    Args:64        image_processor (`Pix2StructImageProcessor`):65            An instance of [`Pix2StructImageProcessor`]. The image processor is a required input.66        tokenizer (Union[`T5TokenizerFast`, `T5Tokenizer`]):67            An instance of ['T5TokenizerFast`] or ['T5Tokenizer`]. The tokenizer is a required input.68    """69 70    attributes = ["image_processor", "tokenizer"]71    image_processor_class = "Pix2StructImageProcessor"72    tokenizer_class = ("T5Tokenizer", "T5TokenizerFast")73 74    def __init__(self, image_processor, tokenizer):75        tokenizer.return_token_type_ids = False76        super().__init__(image_processor, tokenizer)77 78    def __call__(79        self,80        images=None,81        text: Union[TextInput, PreTokenizedInput, list[TextInput], list[PreTokenizedInput]] = None,82        audio=None,83        videos=None,84        **kwargs: Unpack[Pix2StructProcessorKwargs],85    ) -> Union[BatchEncoding, BatchFeature]:86        """87        This method uses [`Pix2StructImageProcessor.preprocess`] method to prepare image(s) for the model, and88        [`T5TokenizerFast.__call__`] to prepare text for the model.89 90        Please refer to the docstring of the above two methods for more information.91        """92        if images is None and text is None:93            raise ValueError("You have to specify either images or text.")94 95        output_kwargs = self._merge_kwargs(96            Pix2StructProcessorKwargs,97            tokenizer_init_kwargs=self.tokenizer.init_kwargs,98            **kwargs,99        )100        add_special_tokens = output_kwargs["text_kwargs"].pop("add_special_tokens", None)101        # Get only text102        if images is None and not self.image_processor.is_vqa:103            output_kwargs["text_kwargs"]["add_special_tokens"] = (104                add_special_tokens if add_special_tokens is not None else True105            )106            self.current_processor = self.tokenizer107            text_encoding = self.tokenizer(text=text, **output_kwargs["text_kwargs"])108            return text_encoding109 110        if not self.image_processor.is_vqa:111            # add pixel_values112            encoding_image_processor = self.image_processor(images, **output_kwargs["images_kwargs"])113        else:114            # add pixel_values and bbox115            output_kwargs["images_kwargs"].setdefault("header_text", text)116            encoding_image_processor = self.image_processor(images, **output_kwargs["images_kwargs"])117 118        if text is not None and not self.image_processor.is_vqa:119            output_kwargs["text_kwargs"]["add_special_tokens"] = (120                add_special_tokens if add_special_tokens is not None else False121            )122            text_encoding = self.tokenizer(text=text, **output_kwargs["text_kwargs"])123 124            if "attention_mask" in text_encoding:125                text_encoding["decoder_attention_mask"] = text_encoding.pop("attention_mask")126            if "input_ids" in text_encoding:127                text_encoding["decoder_input_ids"] = text_encoding.pop("input_ids")128        else:129            text_encoding = None130 131        if text_encoding is not None:132            encoding_image_processor.update(text_encoding)133 134        return encoding_image_processor135 136    @property137    def model_input_names(self):138        image_processor_input_names = self.image_processor.model_input_names139        decoder_ids = ["decoder_attention_mask", "decoder_input_ids"]140        return image_processor_input_names + decoder_ids141 142 143__all__ = ["Pix2StructProcessor"]144