


"""simple docstring"""

import io
import math
from typing import Dict, Optional, Union

import numpy as np
from huggingface_hub import hf_hub_download

from ...image_processing_utils import BaseImageProcessor, BatchFeature
from ...image_transforms import convert_to_rgb, normalize, to_channel_dimension_format, to_pil_image
from ...image_utils import (
    ChannelDimension,
    ImageInput,
    get_image_size,
    infer_channel_dimension_format,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_torch_available, is_vision_available, logging
from ...utils.import_utils import requires_backends


if is_vision_available():
       import textwrap

       from PIL import Image, ImageDraw, ImageFont

if is_torch_available():
       import torch

       from transformers.pytorch_utils import is_torch_greater_or_equal_than_1_11
else:
       UpperCAmelCase__  							=False


UpperCAmelCase__  							=logging.get_logger(__name__)
UpperCAmelCase__  							="ybelkada/fonts"


def lowerCAmelCase_      (				):






 """simple docstring"""







 if is_torch_available() and not is_torch_greater_or_equal_than_1_11:
  raise ImportError(
      f'''You are using torch=={torch.__version__}, but torch>=1.11.0 is required to use '''
      """Pix2StructImageProcessor. Please upgrade torch.""" )


def lowerCAmelCase_      (				UpperCamelCase__	: Optional[int]				,   UpperCamelCase__	: List[str]				,   UpperCamelCase__	: Tuple ):






 """simple docstring"""







 requires_backends(UpperCamelCase__				,   ["""torch"""] )
 _check_torch_version()

 __lowercase									=	image_tensor.unsqueeze(0 )
 __lowercase									=	torch.nn.functional.unfold(UpperCamelCase__				,   (patch_height, patch_width)				,   stride=(patch_height, patch_width) )
 __lowercase									=	patches.reshape(image_tensor.size(0 )				,   image_tensor.size(1 )				,   UpperCamelCase__				,   UpperCamelCase__				,   -1 )
 __lowercase									=	patches.permute(0				,   4				,   2				,   3				,   1 ).reshape(
     image_tensor.size(2 ) // patch_height				,   image_tensor.size(3 ) // patch_width				,   image_tensor.size(1 ) * patch_height * patch_width				,   )
 return patches.unsqueeze(0 )


def lowerCAmelCase_      (				UpperCamelCase__	: str				,   UpperCamelCase__	: int = 36				,   UpperCamelCase__	: str = "black"				,   UpperCamelCase__	: str = "white"				,   UpperCamelCase__	: int = 5				,   UpperCamelCase__	: int = 5				,   UpperCamelCase__	: int = 5				,   UpperCamelCase__	: int = 5				,   UpperCamelCase__	: Optional[bytes] = None				,   UpperCamelCase__	: Optional[str] = None				,   ):






 """simple docstring"""







 requires_backends(UpperCamelCase__				,   """vision""" )
 # Add new lines so that each line is no more than 80 characters.

 __lowercase									=	textwrap.TextWrapper(width=80 )
 __lowercase									=	wrapper.wrap(text=UpperCamelCase__ )
 __lowercase									=	"""\n""".join(UpperCamelCase__ )

 if font_bytes is not None and font_path is None:
  __lowercase									=	io.BytesIO(UpperCamelCase__ )
 elif font_path is not None:
  __lowercase									=	font_path
 else:
  __lowercase									=	hf_hub_download(UpperCamelCase__				,   """Arial.TTF""" )
 __lowercase									=	ImageFont.truetype(UpperCamelCase__				,   encoding="""UTF-8"""				,   size=UpperCamelCase__ )

 # Use a temporary canvas to determine the width and height in pixels when
 # rendering the text.
 __lowercase									=	ImageDraw.Draw(Image.new("""RGB"""				,   (1, 1)				,   UpperCamelCase__ ) )
 __lowercase							,		__lowercase							,		__lowercase							,		__lowercase									=	temp_draw.textbbox((0, 0)				,   UpperCamelCase__				,   UpperCamelCase__ )

 # Create the actual image with a bit of padding around the text.
 __lowercase									=	text_width + left_padding + right_padding
 __lowercase									=	text_height + top_padding + bottom_padding
 __lowercase									=	Image.new("""RGB"""				,   (image_width, image_height)				,   UpperCamelCase__ )
 __lowercase									=	ImageDraw.Draw(UpperCamelCase__ )
 draw.text(xy=(left_padding, top_padding)				,   text=UpperCamelCase__				,   fill=UpperCamelCase__				,   font=UpperCamelCase__ )
 return image


def lowerCAmelCase_      (				UpperCamelCase__	: np.ndarray				,   UpperCamelCase__	: str				,   **UpperCamelCase__	: Any ):






 """simple docstring"""







 requires_backends(UpperCamelCase__				,   """vision""" )

 # Convert to PIL image if necessary
 __lowercase									=	to_pil_image(UpperCamelCase__ )

 __lowercase									=	render_text(UpperCamelCase__				,   **UpperCamelCase__ )
 __lowercase									=	max(header_image.width				,   image.width )

 __lowercase									=	int(image.height * (new_width / image.width) )
 __lowercase									=	int(header_image.height * (new_width / header_image.width) )

 __lowercase									=	Image.new("""RGB"""				,   (new_width, new_height + new_header_height)				,   """white""" )
 new_image.paste(header_image.resize((new_width, new_header_height) )				,   (0, 0) )
 new_image.paste(image.resize((new_width, new_height) )				,   (0, new_header_height) )

 # Convert back to the original framework if necessary
 __lowercase									=	to_numpy_array(UpperCamelCase__ )

 if infer_channel_dimension_format(UpperCamelCase__ ) == ChannelDimension.LAST:
  __lowercase									=	to_channel_dimension_format(UpperCamelCase__				,   ChannelDimension.LAST )

 return new_image





class    lowerCamelCase__						(		_a    ):
      a :     int											=		["""flattened_patches"""]
      def __init__(		self	:	int      ,       A_	:	bool = True      ,       A_	:	bool = True      ,       A_	:	Dict[str, int] = None      ,       A_	:	int = 2_0_4_8      ,       A_	:	bool = False      ,       **A_	:	Optional[int]      ,       ):






       '''simple docstring'''


       super().__init__(**A_							)
       __lowercase									=	patch_size if patch_size is not None else {"""height""": 1_6, """width""": 1_6}
       __lowercase									=	do_normalize
       __lowercase									=	do_convert_rgb
       __lowercase									=	max_patches
       __lowercase									=	is_vqa
      def SCREAMING_SNAKE_CASE_	(		self	:	int      ,       A_	:	np.ndarray      ,       A_	:	int      ,       A_	:	dict      ,       **A_	:	List[Any]							):






       '''simple docstring'''


       requires_backends(self.extract_flattened_patches      ,       """torch"""							)
       _check_torch_version()

       # convert to torch
       __lowercase									=	to_channel_dimension_format(A_      ,       ChannelDimension.FIRST							)
       __lowercase									=	torch.from_numpy(A_							)

       __lowercase							,		__lowercase									=	patch_size["""height"""], patch_size["""width"""]
       __lowercase							,		__lowercase									=	get_image_size(A_							)

       # maximize scale s.t.
       __lowercase									=	math.sqrt(max_patches * (patch_height / image_height) * (patch_width / image_width)							)
       __lowercase									=	max(min(math.floor(scale * image_height / patch_height							)      ,       A_							)      ,       1							)
       __lowercase									=	max(min(math.floor(scale * image_width / patch_width							)      ,       A_							)      ,       1							)
       __lowercase									=	max(num_feasible_rows * patch_height      ,       1							)
       __lowercase									=	max(num_feasible_cols * patch_width      ,       1							)

       __lowercase									=	torch.nn.functional.interpolate(
           image.unsqueeze(0							)      ,       size=(resized_height, resized_width)      ,       mode="""bilinear"""      ,       align_corners=A_      ,       antialias=A_      ,       ).squeeze(0							)

       # [1, rows, columns, patch_height * patch_width * image_channels]
       __lowercase									=	torch_extract_patches(A_      ,       A_      ,       A_							)

       __lowercase									=	patches.shape
       __lowercase									=	patches_shape[1]
       __lowercase									=	patches_shape[2]
       __lowercase									=	patches_shape[3]

       # [rows * columns, patch_height * patch_width * image_channels]
       __lowercase									=	patches.reshape([rows * columns, depth]							)

       # [rows * columns, 1]
       __lowercase									=	torch.arange(A_							).reshape([rows, 1]							).repeat(1      ,       A_							).reshape([rows * columns, 1]							)
       __lowercase									=	torch.arange(A_							).reshape([1, columns]							).repeat(A_      ,       1							).reshape([rows * columns, 1]							)

       # Offset by 1 so the ids do not contain zeros, which represent padding.
       row_ids += 1
       col_ids += 1

       # Prepare additional patch features.
       # [rows * columns, 1]
       __lowercase									=	row_ids.to(torch.floataa							)
       __lowercase									=	col_ids.to(torch.floataa							)

       # [rows * columns, 2 + patch_height * patch_width * image_channels]
       __lowercase									=	torch.cat([row_ids, col_ids, patches]      ,       -1							)

       # [max_patches, 2 + patch_height * patch_width * image_channels]
       __lowercase									=	torch.nn.functional.pad(A_      ,       [0, 0, 0, max_patches - (rows * columns)]							).float()

       __lowercase									=	to_numpy_array(A_							)

       return result
      def SCREAMING_SNAKE_CASE_	(		self	:	str      ,       A_	:	np.ndarray      ,       A_	:	Optional[Union[str, ChannelDimension]] = None      ,       **A_	:	Any							):






       '''simple docstring'''


       if image.dtype == np.uinta:
        __lowercase									=	image.astype(np.floataa							)

       # take mean across the whole `image`
       __lowercase									=	np.mean(A_							)
       __lowercase									=	np.std(A_							)
       __lowercase									=	max(A_      ,       1.0 / math.sqrt(np.prod(image.shape							)							)							)

       return normalize(A_      ,       mean=A_      ,       std=A_      ,       **A_							)






      def SCREAMING_SNAKE_CASE_	(		self	:	int      ,       A_	:	ImageInput      ,       A_	:	Optional[str] = None      ,       A_	:	bool = None      ,       A_	:	Optional[bool] = None      ,       A_	:	Optional[int] = None      ,       A_	:	Optional[Dict[str, int]] = None      ,       A_	:	Optional[Union[str, TensorType]] = None      ,       A_	:	ChannelDimension = ChannelDimension.FIRST      ,       **A_	:	Optional[Any]      ,       ):






       '''simple docstring'''


       __lowercase									=	do_normalize if do_normalize is not None else self.do_normalize
       __lowercase									=	do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb
       __lowercase									=	patch_size if patch_size is not None else self.patch_size
       __lowercase									=	max_patches if max_patches is not None else self.max_patches
       __lowercase									=	self.is_vqa

       if kwargs.get("""data_format"""      ,       A_							) is not None:
        raise ValueError("""data_format is not an accepted input as the outputs are """							)

       __lowercase									=	make_list_of_images(A_							)

       if not valid_images(A_							):
        raise ValueError(
            """Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, """
            """torch.Tensor, tf.Tensor or jax.ndarray."""							)

       # PIL RGBA images are converted to RGB
       if do_convert_rgb:
        __lowercase									=	[convert_to_rgb(A_							) for image in images]

       # All transformations expect numpy arrays.
       __lowercase									=	[to_numpy_array(A_							) for image in images]

       if is_vqa:
        if header_text is None:
         raise ValueError("""A header text must be provided for VQA models."""							)
        __lowercase									=	kwargs.pop("""font_bytes"""      ,       A_							)
        __lowercase									=	kwargs.pop("""font_path"""      ,       A_							)

        if isinstance(A_      ,       A_							):
         __lowercase									=	[header_text] * len(A_							)

        __lowercase									=	[
            render_header(A_      ,       header_text[i]      ,       font_bytes=A_      ,       font_path=A_							)
            for i, image in enumerate(A_							)
        ]

       if do_normalize:
        __lowercase									=	[self.normalize(image=A_							) for image in images]

       # convert to torch tensor and permute
       __lowercase									=	[
           self.extract_flattened_patches(image=A_      ,       max_patches=A_      ,       patch_size=A_							)
           for image in images
       ]

       # create attention mask in numpy
       __lowercase									=	[(image.sum(axis=-1							) != 0).astype(np.floataa							) for image in images]

       __lowercase									=	BatchFeature(
           data={"""flattened_patches""": images, """attention_mask""": attention_masks}      ,       tensor_type=A_							)

       return encoded_outputs




"""simple docstring"""

from collections import OrderedDict

from ...utils import logging
from .auto_factory import _BaseAutoModelClass, _LazyAutoMapping, auto_class_update
from .configuration_auto import CONFIG_MAPPING_NAMES


UpperCAmelCase__  							=logging.get_logger(__name__)


UpperCAmelCase__  							=OrderedDict(
    [
        # Base model mapping
        ("albert", "FlaxAlbertModel"),
        ("bart", "FlaxBartModel"),
        ("beit", "FlaxBeitModel"),
        ("bert", "FlaxBertModel"),
        ("big_bird", "FlaxBigBirdModel"),
        ("blenderbot", "FlaxBlenderbotModel"),
        ("blenderbot-small", "FlaxBlenderbotSmallModel"),
        ("clip", "FlaxCLIPModel"),
        ("distilbert", "FlaxDistilBertModel"),
        ("electra", "FlaxElectraModel"),
        ("gpt-sw3", "FlaxGPT2Model"),
        ("gpt2", "FlaxGPT2Model"),
        ("gpt_neo", "FlaxGPTNeoModel"),
        ("gptj", "FlaxGPTJModel"),
        ("longt5", "FlaxLongT5Model"),
        ("marian", "FlaxMarianModel"),
        ("mbart", "FlaxMBartModel"),
        ("mt5", "FlaxMT5Model"),
        ("opt", "FlaxOPTModel"),
        ("pegasus", "FlaxPegasusModel"),
        ("regnet", "FlaxRegNetModel"),
        ("resnet", "FlaxResNetModel"),
        ("roberta", "FlaxRobertaModel"),
        ("roberta-prelayernorm", "FlaxRobertaPreLayerNormModel"),
        ("roformer", "FlaxRoFormerModel"),
        ("t5", "FlaxT5Model"),
        ("vision-text-dual-encoder", "FlaxVisionTextDualEncoderModel"),
        ("vit", "FlaxViTModel"),
        ("wav2vec2", "FlaxWav2Vec2Model"),
        ("whisper", "FlaxWhisperModel"),
        ("xglm", "FlaxXGLMModel"),
        ("xlm-roberta", "FlaxXLMRobertaModel"),
    ]
)

UpperCAmelCase__  							=OrderedDict(
    [
        # Model for pre-training mapping
        ("albert", "FlaxAlbertForPreTraining"),
        ("bart", "FlaxBartForConditionalGeneration"),
        ("bert", "FlaxBertForPreTraining"),
        ("big_bird", "FlaxBigBirdForPreTraining"),
        ("electra", "FlaxElectraForPreTraining"),
        ("longt5", "FlaxLongT5ForConditionalGeneration"),
        ("mbart", "FlaxMBartForConditionalGeneration"),
        ("mt5", "FlaxMT5ForConditionalGeneration"),
        ("roberta", "FlaxRobertaForMaskedLM"),
        ("roberta-prelayernorm", "FlaxRobertaPreLayerNormForMaskedLM"),
        ("roformer", "FlaxRoFormerForMaskedLM"),
        ("t5", "FlaxT5ForConditionalGeneration"),
        ("wav2vec2", "FlaxWav2Vec2ForPreTraining"),
        ("whisper", "FlaxWhisperForConditionalGeneration"),
        ("xlm-roberta", "FlaxXLMRobertaForMaskedLM"),
    ]
)

UpperCAmelCase__  							=OrderedDict(
    [
        # Model for Masked LM mapping
        ("albert", "FlaxAlbertForMaskedLM"),
        ("bart", "FlaxBartForConditionalGeneration"),
        ("bert", "FlaxBertForMaskedLM"),
        ("big_bird", "FlaxBigBirdForMaskedLM"),
        ("distilbert", "FlaxDistilBertForMaskedLM"),
        ("electra", "FlaxElectraForMaskedLM"),
        ("mbart", "FlaxMBartForConditionalGeneration"),
        ("roberta", "FlaxRobertaForMaskedLM"),
        ("roberta-prelayernorm", "FlaxRobertaPreLayerNormForMaskedLM"),
        ("roformer", "FlaxRoFormerForMaskedLM"),
        ("xlm-roberta", "FlaxXLMRobertaForMaskedLM"),
    ]
)

UpperCAmelCase__  							=OrderedDict(
    [
        # Model for Seq2Seq Causal LM mapping
        ("bart", "FlaxBartForConditionalGeneration"),
        ("blenderbot", "FlaxBlenderbotForConditionalGeneration"),
        ("blenderbot-small", "FlaxBlenderbotSmallForConditionalGeneration"),
        ("encoder-decoder", "FlaxEncoderDecoderModel"),
        ("longt5", "FlaxLongT5ForConditionalGeneration"),
        ("marian", "FlaxMarianMTModel"),
        ("mbart", "FlaxMBartForConditionalGeneration"),
        ("mt5", "FlaxMT5ForConditionalGeneration"),
        ("pegasus", "FlaxPegasusForConditionalGeneration"),
        ("t5", "FlaxT5ForConditionalGeneration"),
    ]
)

UpperCAmelCase__  							=OrderedDict(
    [
        # Model for Image-classsification
        ("beit", "FlaxBeitForImageClassification"),
        ("regnet", "FlaxRegNetForImageClassification"),
        ("resnet", "FlaxResNetForImageClassification"),
        ("vit", "FlaxViTForImageClassification"),
    ]
)

UpperCAmelCase__  							=OrderedDict(
    [
        ("vision-encoder-decoder", "FlaxVisionEncoderDecoderModel"),
    ]
)

UpperCAmelCase__  							=OrderedDict(
    [
        # Model for Causal LM mapping
        ("bart", "FlaxBartForCausalLM"),
        ("bert", "FlaxBertForCausalLM"),
        ("big_bird", "FlaxBigBirdForCausalLM"),
        ("electra", "FlaxElectraForCausalLM"),
        ("gpt-sw3", "FlaxGPT2LMHeadModel"),
        ("gpt2", "FlaxGPT2LMHeadModel"),
        ("gpt_neo", "FlaxGPTNeoForCausalLM"),
        ("gptj", "FlaxGPTJForCausalLM"),
        ("opt", "FlaxOPTForCausalLM"),
        ("roberta", "FlaxRobertaForCausalLM"),
        ("roberta-prelayernorm", "FlaxRobertaPreLayerNormForCausalLM"),
        ("xglm", "FlaxXGLMForCausalLM"),
        ("xlm-roberta", "FlaxXLMRobertaForCausalLM"),
    ]
)

UpperCAmelCase__  							=OrderedDict(
    [
        # Model for Sequence Classification mapping
        ("albert", "FlaxAlbertForSequenceClassification"),
        ("bart", "FlaxBartForSequenceClassification"),
        ("bert", "FlaxBertForSequenceClassification"),
        ("big_bird", "FlaxBigBirdForSequenceClassification"),
        ("distilbert", "FlaxDistilBertForSequenceClassification"),
        ("electra", "FlaxElectraForSequenceClassification"),
        ("mbart", "FlaxMBartForSequenceClassification"),
        ("roberta", "FlaxRobertaForSequenceClassification"),
        ("roberta-prelayernorm", "FlaxRobertaPreLayerNormForSequenceClassification"),
        ("roformer", "FlaxRoFormerForSequenceClassification"),
        ("xlm-roberta", "FlaxXLMRobertaForSequenceClassification"),
    ]
)

UpperCAmelCase__  							=OrderedDict(
    [
        # Model for Question Answering mapping
        ("albert", "FlaxAlbertForQuestionAnswering"),
        ("bart", "FlaxBartForQuestionAnswering"),
        ("bert", "FlaxBertForQuestionAnswering"),
        ("big_bird", "FlaxBigBirdForQuestionAnswering"),
        ("distilbert", "FlaxDistilBertForQuestionAnswering"),
        ("electra", "FlaxElectraForQuestionAnswering"),
        ("mbart", "FlaxMBartForQuestionAnswering"),
        ("roberta", "FlaxRobertaForQuestionAnswering"),
        ("roberta-prelayernorm", "FlaxRobertaPreLayerNormForQuestionAnswering"),
        ("roformer", "FlaxRoFormerForQuestionAnswering"),
        ("xlm-roberta", "FlaxXLMRobertaForQuestionAnswering"),
    ]
)

UpperCAmelCase__  							=OrderedDict(
    [
        # Model for Token Classification mapping
        ("albert", "FlaxAlbertForTokenClassification"),
        ("bert", "FlaxBertForTokenClassification"),
        ("big_bird", "FlaxBigBirdForTokenClassification"),
        ("distilbert", "FlaxDistilBertForTokenClassification"),
        ("electra", "FlaxElectraForTokenClassification"),
        ("roberta", "FlaxRobertaForTokenClassification"),
        ("roberta-prelayernorm", "FlaxRobertaPreLayerNormForTokenClassification"),
        ("roformer", "FlaxRoFormerForTokenClassification"),
        ("xlm-roberta", "FlaxXLMRobertaForTokenClassification"),
    ]
)

UpperCAmelCase__  							=OrderedDict(
    [
        # Model for Multiple Choice mapping
        ("albert", "FlaxAlbertForMultipleChoice"),
        ("bert", "FlaxBertForMultipleChoice"),
        ("big_bird", "FlaxBigBirdForMultipleChoice"),
        ("distilbert", "FlaxDistilBertForMultipleChoice"),
        ("electra", "FlaxElectraForMultipleChoice"),
        ("roberta", "FlaxRobertaForMultipleChoice"),
        ("roberta-prelayernorm", "FlaxRobertaPreLayerNormForMultipleChoice"),
        ("roformer", "FlaxRoFormerForMultipleChoice"),
        ("xlm-roberta", "FlaxXLMRobertaForMultipleChoice"),
    ]
)

UpperCAmelCase__  							=OrderedDict(
    [
        ("bert", "FlaxBertForNextSentencePrediction"),
    ]
)

UpperCAmelCase__  							=OrderedDict(
    [
        ("speech-encoder-decoder", "FlaxSpeechEncoderDecoderModel"),
        ("whisper", "FlaxWhisperForConditionalGeneration"),
    ]
)

UpperCAmelCase__  							=OrderedDict(
    [
        ("whisper", "FlaxWhisperForAudioClassification"),
    ]
)

UpperCAmelCase__  							=_LazyAutoMapping(CONFIG_MAPPING_NAMES, FLAX_MODEL_MAPPING_NAMES)
UpperCAmelCase__  							=_LazyAutoMapping(CONFIG_MAPPING_NAMES, FLAX_MODEL_FOR_PRETRAINING_MAPPING_NAMES)
UpperCAmelCase__  							=_LazyAutoMapping(CONFIG_MAPPING_NAMES, FLAX_MODEL_FOR_MASKED_LM_MAPPING_NAMES)
UpperCAmelCase__  							=_LazyAutoMapping(
    CONFIG_MAPPING_NAMES, FLAX_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES
)
UpperCAmelCase__  							=_LazyAutoMapping(
    CONFIG_MAPPING_NAMES, FLAX_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING_NAMES
)
UpperCAmelCase__  							=_LazyAutoMapping(CONFIG_MAPPING_NAMES, FLAX_MODEL_FOR_VISION_2_SEQ_MAPPING_NAMES)
UpperCAmelCase__  							=_LazyAutoMapping(CONFIG_MAPPING_NAMES, FLAX_MODEL_FOR_CAUSAL_LM_MAPPING_NAMES)
UpperCAmelCase__  							=_LazyAutoMapping(
    CONFIG_MAPPING_NAMES, FLAX_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING_NAMES
)
UpperCAmelCase__  							=_LazyAutoMapping(
    CONFIG_MAPPING_NAMES, FLAX_MODEL_FOR_QUESTION_ANSWERING_MAPPING_NAMES
)
UpperCAmelCase__  							=_LazyAutoMapping(
    CONFIG_MAPPING_NAMES, FLAX_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING_NAMES
)
UpperCAmelCase__  							=_LazyAutoMapping(
    CONFIG_MAPPING_NAMES, FLAX_MODEL_FOR_MULTIPLE_CHOICE_MAPPING_NAMES
)
UpperCAmelCase__  							=_LazyAutoMapping(
    CONFIG_MAPPING_NAMES, FLAX_MODEL_FOR_NEXT_SENTENCE_PREDICTION_MAPPING_NAMES
)
UpperCAmelCase__  							=_LazyAutoMapping(
    CONFIG_MAPPING_NAMES, FLAX_MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING_NAMES
)
UpperCAmelCase__  							=_LazyAutoMapping(
    CONFIG_MAPPING_NAMES, FLAX_MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING_NAMES
)





class    lowerCamelCase__						(		_BaseAutoModelClass    ):
      a :     List[Any]											=		FLAX_MODEL_MAPPING






UpperCAmelCase__  							=auto_class_update(FlaxAutoModel)





class    lowerCamelCase__						(		_BaseAutoModelClass    ):
      a :     int											=		FLAX_MODEL_FOR_PRETRAINING_MAPPING






UpperCAmelCase__  							=auto_class_update(FlaxAutoModelForPreTraining, head_doc="pretraining")





class    lowerCamelCase__						(		_BaseAutoModelClass    ):
      a :     Tuple											=		FLAX_MODEL_FOR_CAUSAL_LM_MAPPING






UpperCAmelCase__  							=auto_class_update(FlaxAutoModelForCausalLM, head_doc="causal language modeling")





class    lowerCamelCase__						(		_BaseAutoModelClass    ):
      a :     List[str]											=		FLAX_MODEL_FOR_MASKED_LM_MAPPING






UpperCAmelCase__  							=auto_class_update(FlaxAutoModelForMaskedLM, head_doc="masked language modeling")





class    lowerCamelCase__						(		_BaseAutoModelClass    ):
      a :     str											=		FLAX_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING






UpperCAmelCase__  							=auto_class_update(
    FlaxAutoModelForSeqaSeqLM, head_doc="sequence-to-sequence language modeling", checkpoint_for_example="t5-base"
)





class    lowerCamelCase__						(		_BaseAutoModelClass    ):
      a :     str											=		FLAX_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING






UpperCAmelCase__  							=auto_class_update(
    FlaxAutoModelForSequenceClassification, head_doc="sequence classification"
)





class    lowerCamelCase__						(		_BaseAutoModelClass    ):
      a :     Tuple											=		FLAX_MODEL_FOR_QUESTION_ANSWERING_MAPPING






UpperCAmelCase__  							=auto_class_update(FlaxAutoModelForQuestionAnswering, head_doc="question answering")





class    lowerCamelCase__						(		_BaseAutoModelClass    ):
      a :     Union[str, Any]											=		FLAX_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING






UpperCAmelCase__  							=auto_class_update(
    FlaxAutoModelForTokenClassification, head_doc="token classification"
)





class    lowerCamelCase__						(		_BaseAutoModelClass    ):
      a :     Optional[Any]											=		FLAX_MODEL_FOR_MULTIPLE_CHOICE_MAPPING






UpperCAmelCase__  							=auto_class_update(FlaxAutoModelForMultipleChoice, head_doc="multiple choice")





class    lowerCamelCase__						(		_BaseAutoModelClass    ):
      a :     Optional[int]											=		FLAX_MODEL_FOR_NEXT_SENTENCE_PREDICTION_MAPPING






UpperCAmelCase__  							=auto_class_update(
    FlaxAutoModelForNextSentencePrediction, head_doc="next sentence prediction"
)





class    lowerCamelCase__						(		_BaseAutoModelClass    ):
      a :     Union[str, Any]											=		FLAX_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING






UpperCAmelCase__  							=auto_class_update(
    FlaxAutoModelForImageClassification, head_doc="image classification"
)





class    lowerCamelCase__						(		_BaseAutoModelClass    ):
      a :     Dict											=		FLAX_MODEL_FOR_VISION_2_SEQ_MAPPING






UpperCAmelCase__  							=auto_class_update(FlaxAutoModelForVisionaSeq, head_doc="vision-to-text modeling")





class    lowerCamelCase__						(		_BaseAutoModelClass    ):
      a :     List[str]											=		FLAX_MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING






UpperCAmelCase__  							=auto_class_update(
    FlaxAutoModelForSpeechSeqaSeq, head_doc="sequence-to-sequence speech-to-text modeling"
)
