



import collections.abc
from typing import Optional, Tuple, Union

import torch
import torch.utils.checkpoint
from torch import nn
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss

from ...activations import ACTaFN
from ...modeling_outputs import BaseModelOutputWithNoAttention, ImageClassifierOutputWithNoAttention
from ...modeling_utils import PreTrainedModel
from ...utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging
from .configuration_poolformer import PoolFormerConfig


a_								=       logging.get_logger(__name__)

# General docstring
a_								=       """PoolFormerConfig"""

# Base docstring
a_								=       """sail/poolformer_s12"""
a_								=       [1, 512, 7, 7]

# Image classification docstring
a_								=       """sail/poolformer_s12"""
a_								=       """tabby, tabby cat"""

a_								=       [
    """sail/poolformer_s12""",
    # See all PoolFormer models at https://huggingface.co/models?filter=poolformer
]





def 		a__      (		_UpperCamelCase      :			List[Any]			,_UpperCamelCase      :			float = 0.0			,_UpperCamelCase      :			bool = False				):
   if drop_prob == 0.0 or not training:
      return input
   __lowerCamelCase    				=     1 - drop_prob
   __lowerCamelCase    				=     (input.shape[0],) + (1,) * (input.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
   __lowerCamelCase    				=     keep_prob + torch.rand(_UpperCamelCase			,dtype=input.dtype			,device=input.device				)
   random_tensor.floor_()  # binarize
   __lowerCamelCase    				=     input.div(_UpperCamelCase				) * random_tensor
   return output



class 			__lowerCAmelCase  (						nn.Module ):




 def __init__(    self						,					__UpperCAmelCase = None					):



    '''simple docstring'''
    super().__init__()
    __lowerCamelCase    				=     drop_prob




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase					):



    '''simple docstring'''
    return drop_path(__UpperCAmelCase						,					self.drop_prob						,					self.training					)




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    return "p={}".format(self.drop_prob					)



class 			__lowerCAmelCase  (						nn.Module ):




 def __init__(    self						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase=None					):



    '''simple docstring'''
    super().__init__()
    __lowerCamelCase    				=     patch_size if isinstance(__UpperCAmelCase						,					collections.abc.Iterable					) else (patch_size, patch_size)
    __lowerCamelCase    				=     stride if isinstance(__UpperCAmelCase						,					collections.abc.Iterable					) else (stride, stride)
    __lowerCamelCase    				=     padding if isinstance(__UpperCAmelCase						,					collections.abc.Iterable					) else (padding, padding)

    __lowerCamelCase    				=     nn.Convad(__UpperCAmelCase						,					__UpperCAmelCase						,					kernel_size=__UpperCAmelCase						,					stride=__UpperCAmelCase						,					padding=__UpperCAmelCase					)
    __lowerCamelCase    				=     norm_layer(__UpperCAmelCase					) if norm_layer else nn.Identity()




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase					):



    '''simple docstring'''
    __lowerCamelCase    				=     self.projection(__UpperCAmelCase					)
    __lowerCamelCase    				=     self.norm(__UpperCAmelCase					)
    return embeddings



class 			__lowerCAmelCase  (						nn.GroupNorm ):




 def __init__(    self						,					__UpperCAmelCase						,					**__UpperCAmelCase					):



    '''simple docstring'''
    super().__init__(1						,					__UpperCAmelCase						,					**__UpperCAmelCase					)



class 			__lowerCAmelCase  (						nn.Module ):




 def __init__(    self						,					__UpperCAmelCase					):



    '''simple docstring'''
    super().__init__()
    __lowerCamelCase    				=     nn.AvgPoolad(__UpperCAmelCase						,					stride=1						,					padding=pool_size // 2						,					count_include_pad=__UpperCAmelCase					)




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase					):



    '''simple docstring'''
    return self.pool(__UpperCAmelCase					) - hidden_states



class 			__lowerCAmelCase  (						nn.Module ):




 def __init__(    self						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase					):



    '''simple docstring'''
    super().__init__()
    __lowerCamelCase    				=     nn.Convad(__UpperCAmelCase						,					__UpperCAmelCase						,					1					)
    __lowerCamelCase    				=     nn.Convad(__UpperCAmelCase						,					__UpperCAmelCase						,					1					)
    __lowerCamelCase    				=     PoolFormerDropPath(__UpperCAmelCase					)
    if isinstance(config.hidden_act						,					__UpperCAmelCase					):
       __lowerCamelCase    				=     ACTaFN[config.hidden_act]
    else:
       __lowerCamelCase    				=     config.hidden_act




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase					):



    '''simple docstring'''
    __lowerCamelCase    				=     self.conva(__UpperCAmelCase					)
    __lowerCamelCase    				=     self.act_fn(__UpperCAmelCase					)
    __lowerCamelCase    				=     self.drop(__UpperCAmelCase					)
    __lowerCamelCase    				=     self.conva(__UpperCAmelCase					)
    __lowerCamelCase    				=     self.drop(__UpperCAmelCase					)

    return hidden_states



class 			__lowerCAmelCase  (						nn.Module ):




 def __init__(    self						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase					):



    '''simple docstring'''
    super().__init__()
    __lowerCamelCase    				=     PoolFormerPooling(__UpperCAmelCase					)
    __lowerCamelCase    				=     PoolFormerOutput(__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase					)
    __lowerCamelCase    				=     PoolFormerGroupNorm(__UpperCAmelCase					)
    __lowerCamelCase    				=     PoolFormerGroupNorm(__UpperCAmelCase					)

    # Useful for training neural nets
    __lowerCamelCase    				=     PoolFormerDropPath(__UpperCAmelCase					) if drop_path > 0.0 else nn.Identity()
    __lowerCamelCase    				=     config.use_layer_scale
    if config.use_layer_scale:
       __lowerCamelCase    				=     nn.Parameter(
           config.layer_scale_init_value * torch.ones((__UpperCAmelCase)					)						,					requires_grad=__UpperCAmelCase					)
       __lowerCamelCase    				=     nn.Parameter(
           config.layer_scale_init_value * torch.ones((__UpperCAmelCase)					)						,					requires_grad=__UpperCAmelCase					)




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase					):



    '''simple docstring'''
    if self.use_layer_scale:
       __lowerCamelCase    				=     self.pooling(self.before_norm(__UpperCAmelCase					)					)
       __lowerCamelCase    				=     self.layer_scale_a.unsqueeze(-1					).unsqueeze(-1					) * pooling_output
       # First residual connection
       __lowerCamelCase    				=     hidden_states + self.drop_path(__UpperCAmelCase					)
       __lowerCamelCase    				=     ()

       __lowerCamelCase    				=     self.output(self.after_norm(__UpperCAmelCase					)					)
       __lowerCamelCase    				=     self.layer_scale_a.unsqueeze(-1					).unsqueeze(-1					) * layer_output
       # Second residual connection
       __lowerCamelCase    				=     hidden_states + self.drop_path(__UpperCAmelCase					)

       __lowerCamelCase    				=     (output,) + outputs
       return outputs

    else:
       __lowerCamelCase    				=     self.drop_path(self.pooling(self.before_norm(__UpperCAmelCase					)					)					)
       # First residual connection
       __lowerCamelCase    				=     pooling_output + hidden_states
       __lowerCamelCase    				=     ()

       # Second residual connection inside the PoolFormerOutput block
       __lowerCamelCase    				=     self.drop_path(self.output(self.after_norm(__UpperCAmelCase					)					)					)
       __lowerCamelCase    				=     hidden_states + layer_output

       __lowerCamelCase    				=     (output,) + outputs
       return outputs



class 			__lowerCAmelCase  (						nn.Module ):




 def __init__(    self						,					__UpperCAmelCase					):



    '''simple docstring'''
    super().__init__()
    __lowerCamelCase    				=     config
    # stochastic depth decay rule
    __lowerCamelCase    				=     [x.item() for x in torch.linspace(0						,					config.drop_path_rate						,					sum(config.depths					)					)]

    # patch embeddings
    __lowerCamelCase    				=     []
    for i in range(config.num_encoder_blocks					):
       embeddings.append(
           PoolFormerEmbeddings(
               patch_size=config.patch_sizes[i]						,					stride=config.strides[i]						,					padding=config.padding[i]						,					num_channels=config.num_channels if i == 0 else config.hidden_sizes[i - 1]						,					hidden_size=config.hidden_sizes[i]						,					)					)
    __lowerCamelCase    				=     nn.ModuleList(__UpperCAmelCase					)

    # Transformer blocks
    __lowerCamelCase    				=     []
    __lowerCamelCase    				=     0
    for i in range(config.num_encoder_blocks					):
       # each block consists of layers
       __lowerCamelCase    				=     []
       if i != 0:
          cur += config.depths[i - 1]
       for j in range(config.depths[i]					):
          layers.append(
              PoolFormerLayer(
                  __UpperCAmelCase						,					num_channels=config.hidden_sizes[i]						,					pool_size=config.pool_size						,					hidden_size=config.hidden_sizes[i]						,					intermediate_size=int(config.hidden_sizes[i] * config.mlp_ratio					)						,					drop_path=dpr[cur + j]						,					)					)
       blocks.append(nn.ModuleList(__UpperCAmelCase					)					)

    __lowerCamelCase    				=     nn.ModuleList(__UpperCAmelCase					)




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase						,					__UpperCAmelCase=False						,					__UpperCAmelCase=True					):



    '''simple docstring'''
    __lowerCamelCase    				=     () if output_hidden_states else None

    __lowerCamelCase    				=     pixel_values
    for idx, layers in enumerate(zip(self.patch_embeddings						,					self.block					)					):
       __lowerCamelCase				,__lowerCamelCase    				=     layers
       # Get patch embeddings from hidden_states
       __lowerCamelCase    				=     embedding_layer(__UpperCAmelCase					)
       # Send the embeddings through the blocks
       for _, blk in enumerate(__UpperCAmelCase					):
          __lowerCamelCase    				=     blk(__UpperCAmelCase					)
          __lowerCamelCase    				=     layer_outputs[0]

       if output_hidden_states:
          __lowerCamelCase    				=     all_hidden_states + (hidden_states,)

    if not return_dict:
       return tuple(v for v in [hidden_states, all_hidden_states] if v is not None					)

    return BaseModelOutputWithNoAttention(last_hidden_state=__UpperCAmelCase						,					hidden_states=__UpperCAmelCase					)



class 			__lowerCAmelCase  (						lowerCAmelCase__ ):
 lowerCAmelCase__							       =	PoolFormerConfig
 lowerCAmelCase__							       =	"""poolformer"""
 lowerCAmelCase__							       =	"""pixel_values"""
 lowerCAmelCase__							       =	True




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase					):



    '''simple docstring'''
    if isinstance(__UpperCAmelCase						,					(nn.Linear, nn.Convad)					):
       module.weight.data.normal_(mean=0.0						,					std=self.config.initializer_range					)
       if module.bias is not None:
          module.bias.data.zero_()
    elif isinstance(__UpperCAmelCase						,					nn.LayerNorm					):
       module.bias.data.zero_()
       module.weight.data.fill_(1.0					)




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase						,					__UpperCAmelCase=False					):



    '''simple docstring'''
    if isinstance(__UpperCAmelCase						,					__UpperCAmelCase					):
       __lowerCamelCase    				=     value



a_								=       R"""
    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use
    it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
    behavior.

    Parameters:
        config ([`PoolFormerConfig`]): Model configuration class with all the parameters of the model.
            Initializing with a config file does not load the weights associated with the model, only the
            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
"""

a_								=       R"""
    Args:
        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See
            [`PoolFormerImageProcessor.__call__`] for details.
"""



@add_start_docstrings(
    """The bare PoolFormer Model transformer outputting raw hidden-states without any specific head on top."""	,						lowerCAmelCase__	,						)
class 			__lowerCAmelCase  (						lowerCAmelCase__ ):




 def __init__(    self						,					__UpperCAmelCase					):



    '''simple docstring'''
    super().__init__(__UpperCAmelCase					)
    __lowerCamelCase    				=     config

    __lowerCamelCase    				=     PoolFormerEncoder(__UpperCAmelCase					)

    # Initialize weights and apply final processing
    self.post_init()




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    return self.embeddings.patch_embeddings




 @add_start_docstrings_to_model_forward(__UpperCAmelCase					)
 @add_code_sample_docstrings(
     checkpoint=_CHECKPOINT_FOR_DOC						,					output_type=__UpperCAmelCase						,					config_class=_CONFIG_FOR_DOC						,					modality='''vision'''						,					expected_output=_EXPECTED_OUTPUT_SHAPE						,					)
 def 							lowerCamelCase   (    self						,					__UpperCAmelCase = None						,					__UpperCAmelCase = None						,					__UpperCAmelCase = None						,					):



    '''simple docstring'''
    __lowerCamelCase    				=     (
        output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
    )
    __lowerCamelCase    				=     return_dict if return_dict is not None else self.config.use_return_dict

    if pixel_values is None:
       raise ValueError('''You have to specify pixel_values'''					)

    __lowerCamelCase    				=     self.encoder(
        __UpperCAmelCase						,					output_hidden_states=__UpperCAmelCase						,					return_dict=__UpperCAmelCase						,					)
    __lowerCamelCase    				=     encoder_outputs[0]

    if not return_dict:
       return (sequence_output, None) + encoder_outputs[1:]

    return BaseModelOutputWithNoAttention(
        last_hidden_state=__UpperCAmelCase						,					hidden_states=encoder_outputs.hidden_states						,					)



class 			__lowerCAmelCase  (						nn.Module ):




 def __init__(    self						,					__UpperCAmelCase					):



    '''simple docstring'''
    super().__init__()
    __lowerCamelCase    				=     nn.Linear(config.hidden_size						,					config.hidden_size					)




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase					):



    '''simple docstring'''
    __lowerCamelCase    				=     self.dense(__UpperCAmelCase					)
    return output



@add_start_docstrings(
    """
    PoolFormer Model transformer with an image classification head on top
    """	,						lowerCAmelCase__	,						)
class 			__lowerCAmelCase  (						lowerCAmelCase__ ):




 def __init__(    self						,					__UpperCAmelCase					):



    '''simple docstring'''
    super().__init__(__UpperCAmelCase					)
    __lowerCamelCase    				=     config.num_labels
    __lowerCamelCase    				=     PoolFormerModel(__UpperCAmelCase					)

    # Final norm
    __lowerCamelCase    				=     PoolFormerGroupNorm(config.hidden_sizes[-1]					)
    # Classifier head
    __lowerCamelCase    				=     (
        nn.Linear(config.hidden_sizes[-1]						,					config.num_labels					) if config.num_labels > 0 else nn.Identity()
    )

    # Initialize weights and apply final processing
    self.post_init()




 @add_start_docstrings_to_model_forward(__UpperCAmelCase					)
 @add_code_sample_docstrings(
     checkpoint=_IMAGE_CLASS_CHECKPOINT						,					output_type=__UpperCAmelCase						,					config_class=_CONFIG_FOR_DOC						,					expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT						,					)
 def 							lowerCamelCase   (    self						,					__UpperCAmelCase = None						,					__UpperCAmelCase = None						,					__UpperCAmelCase = None						,					__UpperCAmelCase = None						,					):



    '''simple docstring'''
    __lowerCamelCase    				=     return_dict if return_dict is not None else self.config.use_return_dict

    __lowerCamelCase    				=     self.poolformer(
        __UpperCAmelCase						,					output_hidden_states=__UpperCAmelCase						,					return_dict=__UpperCAmelCase						,					)

    __lowerCamelCase    				=     outputs[0]

    __lowerCamelCase    				=     self.classifier(self.norm(__UpperCAmelCase					).mean([-2, -1]					)					)

    __lowerCamelCase    				=     None
    if labels is not None:
       if self.config.problem_type is None:
          if self.num_labels == 1:
             __lowerCamelCase    				=     '''regression'''
          elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
             __lowerCamelCase    				=     '''single_label_classification'''
          else:
             __lowerCamelCase    				=     '''multi_label_classification'''

       if self.config.problem_type == "regression":
          __lowerCamelCase    				=     MSELoss()
          if self.num_labels == 1:
             __lowerCamelCase    				=     loss_fct(logits.squeeze()						,					labels.squeeze()					)
          else:
             __lowerCamelCase    				=     loss_fct(__UpperCAmelCase						,					__UpperCAmelCase					)
       elif self.config.problem_type == "single_label_classification":
          __lowerCamelCase    				=     CrossEntropyLoss()
          __lowerCamelCase    				=     loss_fct(logits.view(-1						,					self.num_labels					)						,					labels.view(-1					)					)
       elif self.config.problem_type == "multi_label_classification":
          __lowerCamelCase    				=     BCEWithLogitsLoss()
          __lowerCamelCase    				=     loss_fct(__UpperCAmelCase						,					__UpperCAmelCase					)

    if not return_dict:
       __lowerCamelCase    				=     (logits,) + outputs[2:]
       return ((loss,) + output) if loss is not None else output

    return ImageClassifierOutputWithNoAttention(loss=__UpperCAmelCase						,					logits=__UpperCAmelCase						,					hidden_states=outputs.hidden_states					)











import tempfile
import unittest

from transformers import TaConfig, is_torch_available
from transformers.testing_utils import (
    require_sentencepiece,
    require_tokenizers,
    require_torch,
    slow,
    torch_device,
)

from ...generation.test_utils import GenerationTesterMixin
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch

  from transformers import AutoTokenizer, UMTaForConditionalGeneration, UMTaForQuestionAnswering, UMTaModel



class 			__lowerCAmelCase  :




 def __init__(    self						,					__UpperCAmelCase						,					__UpperCAmelCase=99						,					__UpperCAmelCase=13						,					__UpperCAmelCase=7						,					__UpperCAmelCase=9						,					__UpperCAmelCase=True						,					__UpperCAmelCase=True						,					__UpperCAmelCase=False						,					__UpperCAmelCase=32						,					__UpperCAmelCase=5						,					__UpperCAmelCase=4						,					__UpperCAmelCase=37						,					__UpperCAmelCase=8						,					__UpperCAmelCase=0.1						,					__UpperCAmelCase=0.002						,					__UpperCAmelCase=1						,					__UpperCAmelCase=0						,					__UpperCAmelCase=0						,					__UpperCAmelCase=None						,					__UpperCAmelCase=None						,					):



    '''simple docstring'''
    __lowerCamelCase    				=     parent
    __lowerCamelCase    				=     batch_size
    __lowerCamelCase    				=     encoder_seq_length
    __lowerCamelCase    				=     decoder_seq_length
    # For common tests
    __lowerCamelCase    				=     self.decoder_seq_length
    __lowerCamelCase    				=     is_training
    __lowerCamelCase    				=     use_attention_mask
    __lowerCamelCase    				=     use_labels
    __lowerCamelCase    				=     vocab_size
    __lowerCamelCase    				=     hidden_size
    __lowerCamelCase    				=     num_hidden_layers
    __lowerCamelCase    				=     num_attention_heads
    __lowerCamelCase    				=     d_ff
    __lowerCamelCase    				=     relative_attention_num_buckets
    __lowerCamelCase    				=     dropout_rate
    __lowerCamelCase    				=     initializer_factor
    __lowerCamelCase    				=     eos_token_id
    __lowerCamelCase    				=     pad_token_id
    __lowerCamelCase    				=     decoder_start_token_id
    __lowerCamelCase    				=     None
    __lowerCamelCase    				=     decoder_layers




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    return TaConfig.from_pretrained('''google/umt5-base'''					)




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase=None						,					__UpperCAmelCase=None						,					__UpperCAmelCase=None						,					__UpperCAmelCase=None						,					__UpperCAmelCase=None						,					):



    '''simple docstring'''
    if attention_mask is None:
       __lowerCamelCase    				=     input_ids.ne(config.pad_token_id					)
    if decoder_attention_mask is None:
       __lowerCamelCase    				=     decoder_input_ids.ne(config.pad_token_id					)
    if head_mask is None:
       __lowerCamelCase    				=     torch.ones(config.num_hidden_layers						,					config.num_attention_heads						,					device=__UpperCAmelCase					)
    if decoder_head_mask is None:
       __lowerCamelCase    				=     torch.ones(config.num_decoder_layers						,					config.num_attention_heads						,					device=__UpperCAmelCase					)
    if cross_attn_head_mask is None:
       __lowerCamelCase    				=     torch.ones(
           config.num_decoder_layers						,					config.num_attention_heads						,					device=__UpperCAmelCase					)
    return {
        "input_ids": input_ids,
        "decoder_input_ids": decoder_input_ids,
        "attention_mask": attention_mask,
        "decoder_attention_mask": decoder_attention_mask,
        "head_mask": head_mask,
        "decoder_head_mask": decoder_head_mask,
        "cross_attn_head_mask": cross_attn_head_mask,
    }




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     ids_tensor([self.batch_size, self.encoder_seq_length]						,					self.vocab_size					)
    __lowerCamelCase    				=     ids_tensor([self.batch_size, self.decoder_seq_length]						,					self.vocab_size					)

    # we need to clamp the input ids here to avoid having pad token in between
    # this is because for NllbMoe the position_ids are prepared such that
    # all pad tokens have pos id = 2 and rest are between 2..seq_length
    # and the seq_length here is seq_length - num_pad_tokens
    # but when using past, there is no way of knowing if the past input ids had
    # pad tokens in them, which results in incorrect seq_lenth and which in turn results in
    # position_ids being off by num_pad_tokens in past input
    __lowerCamelCase    				=     input_ids.clamp(self.pad_token_id + 1					)
    __lowerCamelCase    				=     decoder_input_ids.clamp(self.pad_token_id + 1					)

    __lowerCamelCase    				=     self.get_config()
    __lowerCamelCase    				=     config.num_attention_heads
    __lowerCamelCase    				=     self.prepare_inputs_dict(__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase					)
    return config, input_dict




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase				,__lowerCamelCase    				=     self.prepare_config_and_inputs()
    return config, inputs_dict




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    return TaConfig(
        vocab_size=166						,					d_model=self.hidden_size						,					d_ff=self.d_ff						,					d_kv=self.hidden_size // self.num_attention_heads						,					num_layers=self.num_hidden_layers						,					num_decoder_layers=self.decoder_layers						,					num_heads=self.num_attention_heads						,					relative_attention_num_buckets=self.relative_attention_num_buckets						,					dropout_rate=self.dropout_rate						,					initializer_factor=self.initializer_factor						,					eos_token_id=self.eos_token_id						,					bos_token_id=self.pad_token_id						,					pad_token_id=self.pad_token_id						,					decoder_start_token_id=self.decoder_start_token_id						,					)




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    return TaConfig(
        vocab_size=self.vocab_size						,					d_model=self.hidden_size						,					d_ff=self.d_ff						,					d_kv=self.hidden_size // self.num_attention_heads						,					num_layers=self.num_hidden_layers						,					num_decoder_layers=self.decoder_layers						,					num_heads=self.num_attention_heads						,					relative_attention_num_buckets=self.relative_attention_num_buckets						,					dropout_rate=self.dropout_rate						,					initializer_factor=self.initializer_factor						,					eos_token_id=self.eos_token_id						,					bos_token_id=self.pad_token_id						,					pad_token_id=self.pad_token_id						,					decoder_start_token_id=self.decoder_start_token_id						,					)




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					):



    '''simple docstring'''
    __lowerCamelCase    				=     UMTaModel(config=__UpperCAmelCase					)
    model.to(__UpperCAmelCase					)
    model.eval()
    __lowerCamelCase    				=     model(
        input_ids=__UpperCAmelCase						,					decoder_input_ids=__UpperCAmelCase						,					attention_mask=__UpperCAmelCase						,					decoder_attention_mask=__UpperCAmelCase						,					)
    __lowerCamelCase    				=     model(input_ids=__UpperCAmelCase						,					decoder_input_ids=__UpperCAmelCase					)
    __lowerCamelCase    				=     result.last_hidden_state
    __lowerCamelCase    				=     result.past_key_values
    __lowerCamelCase    				=     result.encoder_last_hidden_state

    self.parent.assertEqual(encoder_output.size()						,					(self.batch_size, self.encoder_seq_length, self.hidden_size)					)
    self.parent.assertEqual(decoder_output.size()						,					(self.batch_size, self.decoder_seq_length, self.hidden_size)					)
    # There should be `num_layers` key value embeddings stored in decoder_past
    self.parent.assertEqual(len(__UpperCAmelCase					)						,					config.num_layers					)
    # There should be a self attn key, a self attn value, a cross attn key and a cross attn value stored in each decoder_past tuple
    self.parent.assertEqual(len(decoder_past[0]					)						,					4					)




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					):



    '''simple docstring'''
    __lowerCamelCase    				=     UMTaModel(config=__UpperCAmelCase					).get_decoder().to(__UpperCAmelCase					).eval()
    # first forward pass
    __lowerCamelCase    				=     model(__UpperCAmelCase						,					use_cache=__UpperCAmelCase					)
    __lowerCamelCase    				=     model(__UpperCAmelCase					)
    __lowerCamelCase    				=     model(__UpperCAmelCase						,					use_cache=__UpperCAmelCase					)

    self.parent.assertTrue(len(__UpperCAmelCase					) == len(__UpperCAmelCase					)					)
    self.parent.assertTrue(len(__UpperCAmelCase					) == len(__UpperCAmelCase					) + 1					)

    __lowerCamelCase				,__lowerCamelCase    				=     outputs.to_tuple()

    # create hypothetical next token and extent to next_input_ids
    __lowerCamelCase    				=     ids_tensor((self.batch_size, 1)						,					config.vocab_size					)

    # append to next input_ids and
    __lowerCamelCase    				=     torch.cat([input_ids, next_tokens]						,					dim=-1					)

    __lowerCamelCase    				=     model(__UpperCAmelCase					)['''last_hidden_state''']
    __lowerCamelCase    				=     model(__UpperCAmelCase						,					past_key_values=__UpperCAmelCase					)['''last_hidden_state''']

    # select random slice
    __lowerCamelCase    				=     ids_tensor((1,)						,					output_from_past.shape[-1]					).item()
    __lowerCamelCase    				=     output_from_no_past[:, -1, random_slice_idx].detach()
    __lowerCamelCase    				=     output_from_past[:, 0, random_slice_idx].detach()

    # test that outputs are equal for slice
    self.parent.assertTrue(torch.allclose(__UpperCAmelCase						,					__UpperCAmelCase						,					atol=1E-3					)					)




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase						,					__UpperCAmelCase						,					):



    '''simple docstring'''
    __lowerCamelCase    				=     UMTaModel(config=__UpperCAmelCase					).to(__UpperCAmelCase					).half().eval()
    __lowerCamelCase    				=     model(**__UpperCAmelCase					)['''last_hidden_state''']
    self.parent.assertFalse(torch.isnan(__UpperCAmelCase					).any().item()					)



@require_torch
class 			__lowerCAmelCase  (						lowerCAmelCase__	,						lowerCAmelCase__	,						lowerCAmelCase__	,						unittest.TestCase ):
 lowerCAmelCase__							       =	(
     (UMTaModel, UMTaForConditionalGeneration, UMTaForQuestionAnswering) if is_torch_available() else ()
 )
 lowerCAmelCase__							       =	(UMTaForConditionalGeneration,) if is_torch_available() else ()
 lowerCAmelCase__							       =	(
     {
         """conversational""": UMTaForConditionalGeneration,
         """feature-extraction""": UMTaModel,
         """summarization""": UMTaForConditionalGeneration,
         """text2text-generation""": UMTaForConditionalGeneration,
         """translation""": UMTaForConditionalGeneration,
         """question-answering""": UMTaForQuestionAnswering,
     }
     if is_torch_available()
     else {}
 )
 lowerCAmelCase__							       =	True
 lowerCAmelCase__							       =	False
 lowerCAmelCase__							       =	False
 lowerCAmelCase__							       =	True
 lowerCAmelCase__							       =	True
 # The small UMT5 model needs higher percentages for CPU/MP tests
 lowerCAmelCase__							       =	[0.8, 0.9]




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     UMTaModelTester(self					)




 @unittest.skip('''Test has a segmentation fault on torch 1.8.0'''					)
 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     self.model_tester.prepare_config_and_inputs()
    __lowerCamelCase    				=     UMTaModel(config_and_inputs[0]					).to(__UpperCAmelCase					)
    with tempfile.TemporaryDirectory() as tmpdirname:
       torch.onnx.export(
           __UpperCAmelCase						,					(config_and_inputs[1], config_and_inputs[3], config_and_inputs[2])						,					F"""{tmpdirname}/t5_test.onnx"""						,					export_params=__UpperCAmelCase						,					opset_version=9						,					input_names=['''input_ids''', '''decoder_input_ids''']						,					)




 @unittest.skipIf(torch_device == '''cpu'''						,					'''Cant do half precision'''					)
 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_model_fpaa_forward(*__UpperCAmelCase					)




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     ['''encoder_attentions''', '''decoder_attentions''', '''cross_attentions''']
    __lowerCamelCase    				=     self.model_tester.prepare_config_and_inputs()
    __lowerCamelCase    				=     config_and_inputs[0]
    __lowerCamelCase    				=     UMTaForConditionalGeneration(__UpperCAmelCase					).eval()
    model.to(__UpperCAmelCase					)

    __lowerCamelCase    				=     {
        '''head_mask''': torch.zeros(config.num_layers						,					config.num_heads						,					device=__UpperCAmelCase					),
        '''decoder_head_mask''': torch.zeros(config.num_decoder_layers						,					config.num_heads						,					device=__UpperCAmelCase					),
        '''cross_attn_head_mask''': torch.zeros(config.num_decoder_layers						,					config.num_heads						,					device=__UpperCAmelCase					),
    }

    for attn_name, (name, mask) in zip(__UpperCAmelCase						,					head_masking.items()					):
       __lowerCamelCase    				=     {name: mask}
       # Explicitly pass decoder_head_mask as it is required from T5 model when head_mask specified
       if name == "head_mask":
          __lowerCamelCase    				=     torch.ones(
              config.num_decoder_layers						,					config.num_heads						,					device=__UpperCAmelCase					)

       __lowerCamelCase    				=     model.generate(
           config_and_inputs[1]['''input_ids''']						,					num_beams=1						,					max_length=3						,					output_attentions=__UpperCAmelCase						,					return_dict_in_generate=__UpperCAmelCase						,					**__UpperCAmelCase						,					)
       # We check the state of decoder_attentions and cross_attentions just from the last step
       __lowerCamelCase    				=     out[attn_name] if attn_name == attention_names[0] else out[attn_name][-1]
       self.assertEqual(sum([w.sum().item() for w in attn_weights]					)						,					0.0					)




 @unittest.skip('''Does not work on the tiny model as we keep hitting edge cases.'''					)
 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    pass



@require_torch
@require_sentencepiece
@require_tokenizers
class 			__lowerCAmelCase  (						unittest.TestCase ):




 @slow
 @unittest.skip(
     '''Unless we stop stripping left and right by default for all special tokens, the expected ids obtained here will not match the original ones. Wait for https://github.com/huggingface/transformers/pull/23909 to be merged'''					)
 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     UMTaForConditionalGeneration.from_pretrained('''google/umt5-small'''						,					return_dict=__UpperCAmelCase					).to(__UpperCAmelCase					)
    __lowerCamelCase    				=     AutoTokenizer.from_pretrained('''google/umt5-small'''						,					use_fast=__UpperCAmelCase						,					legacy=__UpperCAmelCase					)
    __lowerCamelCase    				=     [
        '''Bonjour monsieur <extra_id_0> bien <extra_id_1>.''',
        '''No se como puedo <extra_id_0>.''',
        '''This is the reason why we <extra_id_0> them.''',
        '''The <extra_id_0> walks in <extra_id_1>, seats''',
        '''A <extra_id_0> walks into a bar and orders a <extra_id_1> with <extra_id_2> pinch of <extra_id_3>.''',
    ]
    __lowerCamelCase    				=     tokenizer(__UpperCAmelCase						,					return_tensors='''pt'''						,					padding=__UpperCAmelCase					).input_ids
    # fmt: off
    __lowerCamelCase    				=     torch.tensor(
        [
            [ 38530, 210703, 256299, 1410, 256298, 274, 1, 0,0, 0, 0, 0, 0, 0, 0, 0,0, 0],
            [   826, 321, 671, 25922, 256299, 274, 1, 0,0, 0, 0, 0, 0, 0, 0, 0,0, 0],
            [  1460, 339, 312, 19014, 10620, 758, 256299, 2355,274, 1, 0, 0, 0, 0, 0, 0,0, 0],
            [   517, 256299, 14869, 281, 301, 256298, 275, 119983,1, 0, 0, 0, 0, 0, 0, 0,0, 0],
            [   320, 256299, 14869, 281, 2234, 289, 2275, 333,61391, 289, 256298, 543, 256297, 168714, 329, 256296,274, 1],
        ]					)
    # fmt: on
    torch.testing.assert_allclose(__UpperCAmelCase						,					__UpperCAmelCase					)

    __lowerCamelCase    				=     model.generate(input_ids.to(__UpperCAmelCase					)					)
    __lowerCamelCase    				=     [
        '''<pad><extra_id_0> et<extra_id_1> [eod] <extra_id_2><extra_id_55>.. [eod] 💐 💐 💐 💐 💐 💐 💐 💐 💐 💐 💐 <extra_id_56>ajšietosto<extra_id_56>lleux<extra_id_19><extra_id_6>ajšie</s>''',
        '''<pad><extra_id_0>.<extra_id_1>.,<0x0A>...spech <0x0A><extra_id_20> <extra_id_21></s><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>''',
        '''<pad><extra_id_0> are not going to be a part of the world. We are not going to be a part of<extra_id_1> and<extra_id_2><0x0A><extra_id_48>.<extra_id_48></s><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>''',
        '''<pad><extra_id_0> door<extra_id_1>, the door<extra_id_2> 피해[/</s><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>''',
        '''<pad><extra_id_0>nyone who<extra_id_1> drink<extra_id_2> a<extra_id_3> alcohol<extra_id_4> A<extra_id_5> A. This<extra_id_6> I<extra_id_7><extra_id_52><extra_id_53></s><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>''',
    ]
    __lowerCamelCase    				=     tokenizer.batch_decode(__UpperCAmelCase					)
    self.assertEqual(__UpperCAmelCase						,					__UpperCAmelCase					)






