from typing import Optional, Tuple, Union

import tensorflow as tf

from ...activations_tf import ACTaFN
from ...file_utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward
from ...modeling_tf_outputs import (
    TFBaseModelOutputWithNoAttention,
    TFBaseModelOutputWithPoolingAndNoAttention,
    TFSequenceClassifierOutput,
)
from ...modeling_tf_utils import TFPreTrainedModel, TFSequenceClassificationLoss, keras_serializable, unpack_inputs
from ...tf_utils import shape_list
from ...utils import logging
from .configuration_regnet import RegNetConfig


UpperCamelCase__    							= logging.get_logger(__name__)

# General docstring
UpperCamelCase__    							= 'RegNetConfig'

# Base docstring
UpperCamelCase__    							= 'facebook/regnet-y-040'
UpperCamelCase__    							= [1, 1_0_8_8, 7, 7]

# Image classification docstring
UpperCamelCase__    							= 'facebook/regnet-y-040'
UpperCamelCase__    							= 'tabby, tabby cat'

UpperCamelCase__    							= [
    'facebook/regnet-y-040',
    # See all regnet models at https://huggingface.co/models?filter=regnet
]





class 					A							(    tf.keras.layers.Layer					):


			def __init__(self					:						List[str]	,  __UpperCAmelCase					:						int	,  __UpperCAmelCase					:						int = 3	,  __UpperCAmelCase					:						int = 1	,  __UpperCAmelCase					:						int = 1	,  __UpperCAmelCase					:						Optional[str] = "relu"	,  **__UpperCAmelCase					:						Union[str, Any]	,  )						->    List[Any]:
									"""simple docstring"""



									super().__init__(**_lowerCAmelCase		)
									# The padding and conv has been verified in
									# https://colab.research.google.com/gist/sayakpaul/854bc10eeaf21c9ee2119e0b9f3841a7/scratchpad.ipynb
									UpperCAmelCase__								=  tf.keras.layers.ZeroPaddingaD(padding=kernel_size // 2		)
									UpperCAmelCase__								=  tf.keras.layers.ConvaD(
									    filters=_lowerCAmelCase	,  kernel_size=_lowerCAmelCase	,  strides=_lowerCAmelCase	,  padding="VALID"	,  groups=_lowerCAmelCase	,  use_bias=_lowerCAmelCase	,  name="convolution"	,  )
									UpperCAmelCase__								=  tf.keras.layers.BatchNormalization(epsilon=1E-5	,  momentum=0.9	,  name="normalization"		)
									UpperCAmelCase__								=  ACTaFN[activation] if activation is not None else tf.identity





			def       lowercase_				(self					:						Optional[Any]	,  __UpperCAmelCase					:						Dict		)						->    Optional[Any]:
									"""simple docstring"""



									UpperCAmelCase__								=  self.convolution(self.padding(_lowerCAmelCase		)		)
									UpperCAmelCase__								=  self.normalization(_lowerCAmelCase		)
									UpperCAmelCase__								=  self.activation(_lowerCAmelCase		)
									return hidden_state





class 					A							(    tf.keras.layers.Layer					):


			def __init__(self					:						str	,  __UpperCAmelCase					:						RegNetConfig	,  **__UpperCAmelCase					:						List[str]		)						->    List[str]:
									"""simple docstring"""



									super().__init__(**_lowerCAmelCase		)
									UpperCAmelCase__								=  config.num_channels
									UpperCAmelCase__								=  TFRegNetConvLayer(
									    out_channels=config.embedding_size	,  kernel_size=3	,  stride=2	,  activation=config.hidden_act	,  name="embedder"	,  )





			def       lowercase_				(self					:						List[Any]	,  __UpperCAmelCase					:						Optional[int]		)						->    List[str]:
									"""simple docstring"""



									UpperCAmelCase__								=  shape_list(_lowerCAmelCase		)[1]
									if tf.executing_eagerly() and num_channels != self.num_channels:
															raise ValueError(
															    "Make sure that the channel dimension of the pixel values match with the one set in the configuration."		)

									# When running on CPU, `tf.keras.layers.Conv2D` doesn't support `NCHW` format.
									# So change the input format from `NCHW` to `NHWC`.
									# shape = (batch_size, in_height, in_width, in_channels=num_channels)
									UpperCAmelCase__								=  tf.transpose(_lowerCAmelCase	,  perm=(0, 2, 3, 1)		)
									UpperCAmelCase__								=  self.embedder(_lowerCAmelCase		)
									return hidden_state





class 					A							(    tf.keras.layers.Layer					):


			def __init__(self					:						List[Any]	,  __UpperCAmelCase					:						int	,  __UpperCAmelCase					:						int = 2	,  **__UpperCAmelCase					:						Dict		)						->    Dict:
									"""simple docstring"""



									super().__init__(**_lowerCAmelCase		)
									UpperCAmelCase__								=  tf.keras.layers.ConvaD(
									    filters=_lowerCAmelCase	,  kernel_size=1	,  strides=_lowerCAmelCase	,  use_bias=_lowerCAmelCase	,  name="convolution"		)
									UpperCAmelCase__								=  tf.keras.layers.BatchNormalization(epsilon=1E-5	,  momentum=0.9	,  name="normalization"		)





			def       lowercase_				(self					:						List[str]	,  __UpperCAmelCase					:						tf.Tensor	,  __UpperCAmelCase					:						bool = False		)						->    Optional[int]:
									"""simple docstring"""



									return self.normalization(self.convolution(_lowerCAmelCase		)	,  training=_lowerCAmelCase		)





class 					A							(    tf.keras.layers.Layer					):


			def __init__(self					:						List[str]	,  __UpperCAmelCase					:						int	,  __UpperCAmelCase					:						int	,  **__UpperCAmelCase					:						Optional[int]		)						->    str:
									"""simple docstring"""



									super().__init__(**_lowerCAmelCase		)
									UpperCAmelCase__								=  tf.keras.layers.GlobalAveragePoolingaD(keepdims=_lowerCAmelCase	,  name="pooler"		)
									UpperCAmelCase__								=  [
									    tf.keras.layers.ConvaD(filters=_lowerCAmelCase	,  kernel_size=1	,  activation="relu"	,  name="attention.0"		),
									    tf.keras.layers.ConvaD(filters=_lowerCAmelCase	,  kernel_size=1	,  activation="sigmoid"	,  name="attention.2"		),
									]





			def       lowercase_				(self					:						Union[str, Any]	,  __UpperCAmelCase					:						Optional[int]		)						->    List[str]:
									"""simple docstring"""



									UpperCAmelCase__								=  self.pooler(_lowerCAmelCase		)
									for layer_module in self.attention:
															UpperCAmelCase__								=  layer_module(_lowerCAmelCase		)
									UpperCAmelCase__								=  hidden_state * pooled
									return hidden_state





class 					A							(    tf.keras.layers.Layer					):


			def __init__(self					:						str	,  __UpperCAmelCase					:						RegNetConfig	,  __UpperCAmelCase					:						int	,  __UpperCAmelCase					:						int	,  __UpperCAmelCase					:						int = 1	,  **__UpperCAmelCase					:						Dict		)						->    List[Any]:
									"""simple docstring"""



									super().__init__(**_lowerCAmelCase		)
									UpperCAmelCase__								=  in_channels != out_channels or stride != 1
									UpperCAmelCase__								=  max(1	,  out_channels // config.groups_width		)
									UpperCAmelCase__								=  (
									    TFRegNetShortCut(_lowerCAmelCase	,  stride=_lowerCAmelCase	,  name="shortcut"		)
									    if should_apply_shortcut
									    else tf.keras.layers.Activation("linear"	,  name="shortcut"		)
									)
									# `self.layers` instead of `self.layer` because that is a reserved argument.
									UpperCAmelCase__								=  [
									    TFRegNetConvLayer(_lowerCAmelCase	,  kernel_size=1	,  activation=config.hidden_act	,  name="layer.0"		),
									    TFRegNetConvLayer(
									        _lowerCAmelCase	,  stride=_lowerCAmelCase	,  groups=_lowerCAmelCase	,  activation=config.hidden_act	,  name="layer.1"		),
									    TFRegNetConvLayer(_lowerCAmelCase	,  kernel_size=1	,  activation=_lowerCAmelCase	,  name="layer.2"		),
									]
									UpperCAmelCase__								=  ACTaFN[config.hidden_act]





			def       lowercase_				(self					:						Tuple	,  __UpperCAmelCase					:						Union[str, Any]		)						->    Union[str, Any]:
									"""simple docstring"""



									UpperCAmelCase__								=  hidden_state
									for layer_module in self.layers:
															UpperCAmelCase__								=  layer_module(_lowerCAmelCase		)
									UpperCAmelCase__								=  self.shortcut(_lowerCAmelCase		)
									hidden_state += residual
									UpperCAmelCase__								=  self.activation(_lowerCAmelCase		)
									return hidden_state





class 					A							(    tf.keras.layers.Layer					):


			def __init__(self					:						Optional[int]	,  __UpperCAmelCase					:						RegNetConfig	,  __UpperCAmelCase					:						int	,  __UpperCAmelCase					:						int	,  __UpperCAmelCase					:						int = 1	,  **__UpperCAmelCase					:						Optional[int]		)						->    Tuple:
									"""simple docstring"""



									super().__init__(**_lowerCAmelCase		)
									UpperCAmelCase__								=  in_channels != out_channels or stride != 1
									UpperCAmelCase__								=  max(1	,  out_channels // config.groups_width		)
									UpperCAmelCase__								=  (
									    TFRegNetShortCut(_lowerCAmelCase	,  stride=_lowerCAmelCase	,  name="shortcut"		)
									    if should_apply_shortcut
									    else tf.keras.layers.Activation("linear"	,  name="shortcut"		)
									)
									UpperCAmelCase__								=  [
									    TFRegNetConvLayer(_lowerCAmelCase	,  kernel_size=1	,  activation=config.hidden_act	,  name="layer.0"		),
									    TFRegNetConvLayer(
									        _lowerCAmelCase	,  stride=_lowerCAmelCase	,  groups=_lowerCAmelCase	,  activation=config.hidden_act	,  name="layer.1"		),
									    TFRegNetSELayer(_lowerCAmelCase	,  reduced_channels=int(round(in_channels / 4		)		)	,  name="layer.2"		),
									    TFRegNetConvLayer(_lowerCAmelCase	,  kernel_size=1	,  activation=_lowerCAmelCase	,  name="layer.3"		),
									]
									UpperCAmelCase__								=  ACTaFN[config.hidden_act]





			def       lowercase_				(self					:						Union[str, Any]	,  __UpperCAmelCase					:						List[Any]		)						->    str:
									"""simple docstring"""



									UpperCAmelCase__								=  hidden_state
									for layer_module in self.layers:
															UpperCAmelCase__								=  layer_module(_lowerCAmelCase		)
									UpperCAmelCase__								=  self.shortcut(_lowerCAmelCase		)
									hidden_state += residual
									UpperCAmelCase__								=  self.activation(_lowerCAmelCase		)
									return hidden_state





class 					A							(    tf.keras.layers.Layer					):


			def __init__(self					:						str	,  __UpperCAmelCase					:						RegNetConfig	,  __UpperCAmelCase					:						int	,  __UpperCAmelCase					:						int	,  __UpperCAmelCase					:						int = 2	,  __UpperCAmelCase					:						int = 2	,  **__UpperCAmelCase					:						List[Any]		)						->    Optional[Any]:
									"""simple docstring"""



									super().__init__(**_lowerCAmelCase		)

									UpperCAmelCase__								=  TFRegNetXLayer if config.layer_type == """x""" else TFRegNetYLayer
									UpperCAmelCase__								=  [
									    # downsampling is done in the first layer with stride of 2
									    layer(_lowerCAmelCase	,  _lowerCAmelCase	,  _lowerCAmelCase	,  stride=_lowerCAmelCase	,  name="layers.0"		),
									    *[layer(_lowerCAmelCase	,  _lowerCAmelCase	,  _lowerCAmelCase	,  name=f"""layers.{i+1}"""		) for i in range(depth - 1		)],
									]





			def       lowercase_				(self					:						List[Any]	,  __UpperCAmelCase					:						Dict		)						->    Optional[Any]:
									"""simple docstring"""



									for layer_module in self.layers:
															UpperCAmelCase__								=  layer_module(_lowerCAmelCase		)
									return hidden_state





class 					A							(    tf.keras.layers.Layer					):


			def __init__(self					:						Tuple	,  __UpperCAmelCase					:						RegNetConfig	,  **__UpperCAmelCase					:						int		)						->    List[str]:
									"""simple docstring"""



									super().__init__(**_lowerCAmelCase		)
									UpperCAmelCase__								=  []
									# based on `downsample_in_first_stage`, the first layer of the first stage may or may not downsample the input
									self.stages.append(
									    TFRegNetStage(
									        _lowerCAmelCase	,  config.embedding_size	,  config.hidden_sizes[0]	,  stride=2 if config.downsample_in_first_stage else 1	,  depth=config.depths[0]	,  name="stages.0"	,  )		)
									UpperCAmelCase__								=  zip(config.hidden_sizes	,  config.hidden_sizes[1:]		)
									for i, ((in_channels, out_channels), depth) in enumerate(zip(_lowerCAmelCase	,  config.depths[1:]		)		):
															self.stages.append(TFRegNetStage(_lowerCAmelCase	,  _lowerCAmelCase	,  _lowerCAmelCase	,  depth=_lowerCAmelCase	,  name=f"""stages.{i+1}"""		)		)





			def       lowercase_				(self					:						Any	,  __UpperCAmelCase					:						tf.Tensor	,  __UpperCAmelCase					:						bool = False	,  __UpperCAmelCase					:						bool = True		)						->    Optional[int]:
									"""simple docstring"""



									UpperCAmelCase__								=  () if output_hidden_states else None

									for stage_module in self.stages:
															if output_hidden_states:
																					UpperCAmelCase__								=  hidden_states + (hidden_state,)

															UpperCAmelCase__								=  stage_module(_lowerCAmelCase		)

									if output_hidden_states:
															UpperCAmelCase__								=  hidden_states + (hidden_state,)

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

									return TFBaseModelOutputWithNoAttention(last_hidden_state=_lowerCAmelCase	,  hidden_states=_lowerCAmelCase		)





@keras_serializable
class 					A							(    tf.keras.layers.Layer					):
			__UpperCAmelCase			:			List[Any]              =      RegNetConfig


			def __init__(self					:						Tuple	,  __UpperCAmelCase					:						List[Any]	,  **__UpperCAmelCase					:						str		)						->    Optional[int]:
									"""simple docstring"""



									super().__init__(**_lowerCAmelCase		)
									UpperCAmelCase__								=  config
									UpperCAmelCase__								=  TFRegNetEmbeddings(_lowerCAmelCase	,  name="embedder"		)
									UpperCAmelCase__								=  TFRegNetEncoder(_lowerCAmelCase	,  name="encoder"		)
									UpperCAmelCase__								=  tf.keras.layers.GlobalAveragePoolingaD(keepdims=_lowerCAmelCase	,  name="pooler"		)





			@unpack_inputs
			def       lowercase_				(self					:						Dict	,  __UpperCAmelCase					:						tf.Tensor	,  __UpperCAmelCase					:						Optional[bool] = None	,  __UpperCAmelCase					:						Optional[bool] = None	,  __UpperCAmelCase					:						bool = False	,  )						->    Union[str, Any]:
									"""simple docstring"""



									UpperCAmelCase__								=  (
									    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
									)
									UpperCAmelCase__								=  return_dict if return_dict is not None else self.config.use_return_dict

									UpperCAmelCase__								=  self.embedder(_lowerCAmelCase	,  training=_lowerCAmelCase		)

									UpperCAmelCase__								=  self.encoder(
									    _lowerCAmelCase	,  output_hidden_states=_lowerCAmelCase	,  return_dict=_lowerCAmelCase	,  training=_lowerCAmelCase		)

									UpperCAmelCase__								=  encoder_outputs[0]
									UpperCAmelCase__								=  self.pooler(_lowerCAmelCase		)

									# Change to NCHW output format have uniformity in the modules
									UpperCAmelCase__								=  tf.transpose(_lowerCAmelCase	,  perm=(0, 3, 1, 2)		)
									UpperCAmelCase__								=  tf.transpose(_lowerCAmelCase	,  perm=(0, 3, 1, 2)		)

									# Change the other hidden state outputs to NCHW as well
									if output_hidden_states:
															UpperCAmelCase__								=  tuple([tf.transpose(_lowerCAmelCase	,  perm=(0, 3, 1, 2)		) for h in encoder_outputs[1]]		)

									if not return_dict:
															return (last_hidden_state, pooled_output) + encoder_outputs[1:]

									return TFBaseModelOutputWithPoolingAndNoAttention(
									    last_hidden_state=_lowerCAmelCase	,  pooler_output=_lowerCAmelCase	,  hidden_states=hidden_states if output_hidden_states else encoder_outputs.hidden_states	,  )





class 					A							(    __UpperCamelCase					):
			__UpperCAmelCase			:			Union[str, Any]              =      RegNetConfig
			__UpperCAmelCase			:			str              =      "regnet"
			__UpperCAmelCase			:			Any              =      "pixel_values"


			@property
			def       lowercase_				(self					:						Optional[int]		)						->    Optional[int]:
									"""simple docstring"""



									return {"pixel_values": tf.TensorSpec(shape=(None, self.config.num_channels, 2_2_4, 2_2_4)	,  dtype=tf.floataa		)}




UpperCamelCase__    							= R'\n    Parameters:\n    This model is a Tensorflow\n    [tf.keras.layers.Layer](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Layer) sub-class. Use it as a\n    regular Tensorflow Module and refer to the Tensorflow documentation for all matter related to general usage and\n    behavior.\n        config ([`RegNetConfig`]): Model configuration class with all the parameters of the model.\n            Initializing with a config file does not load the weights associated with the model, only the\n            configuration. Check out the [`~TFPreTrainedModel.from_pretrained`] method to load the model weights.\n'

UpperCamelCase__    							= R'\n    Args:\n        pixel_values (`tf.Tensor` of shape `(batch_size, num_channels, height, width)`):\n            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See\n            [`ConveNextImageProcessor.__call__`] for details.\n        output_hidden_states (`bool`, *optional*):\n            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for\n            more detail.\n        return_dict (`bool`, *optional*):\n            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.\n'





@add_start_docstrings(
    'The bare RegNet model outputting raw features without any specific head on top.'		,    __UpperCamelCase		,    )
class 					A							(    __UpperCamelCase					):


			def __init__(self					:						List[Any]	,  __UpperCAmelCase					:						RegNetConfig	,  *__UpperCAmelCase					:						List[str]	,  **__UpperCAmelCase					:						List[str]		)						->    List[Any]:
									"""simple docstring"""



									super().__init__(_lowerCAmelCase	,  *_lowerCAmelCase	,  **_lowerCAmelCase		)
									UpperCAmelCase__								=  TFRegNetMainLayer(_lowerCAmelCase	,  name="regnet"		)





			@unpack_inputs
			@add_start_docstrings_to_model_forward(_lowerCAmelCase		)
			@add_code_sample_docstrings(
			    checkpoint=_CHECKPOINT_FOR_DOC	,  output_type=_lowerCAmelCase	,  config_class=_CONFIG_FOR_DOC	,  modality="vision"	,  expected_output=_EXPECTED_OUTPUT_SHAPE	,  )
			def       lowercase_				(self					:						str	,  __UpperCAmelCase					:						tf.Tensor	,  __UpperCAmelCase					:						Optional[bool] = None	,  __UpperCAmelCase					:						Optional[bool] = None	,  __UpperCAmelCase					:						List[Any]=False	,  )						->    int:
									"""simple docstring"""



									UpperCAmelCase__								=  (
									    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
									)
									UpperCAmelCase__								=  return_dict if return_dict is not None else self.config.use_return_dict

									UpperCAmelCase__								=  self.regnet(
									    pixel_values=_lowerCAmelCase	,  output_hidden_states=_lowerCAmelCase	,  return_dict=_lowerCAmelCase	,  training=_lowerCAmelCase	,  )
									if not return_dict:
															return (outputs[0],) + outputs[1:]

									return TFBaseModelOutputWithPoolingAndNoAttention(
									    last_hidden_state=outputs.last_hidden_state	,  pooler_output=outputs.pooler_output	,  hidden_states=outputs.hidden_states	,  )





@add_start_docstrings(
    '\n    RegNet Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for\n    ImageNet.\n    '		,    __UpperCamelCase		,    )
class 					A							(    __UpperCamelCase		,    __UpperCamelCase					):


			def __init__(self					:						str	,  __UpperCAmelCase					:						RegNetConfig	,  *__UpperCAmelCase					:						Dict	,  **__UpperCAmelCase					:						Optional[Any]		)						->    Optional[Any]:
									"""simple docstring"""



									super().__init__(_lowerCAmelCase	,  *_lowerCAmelCase	,  **_lowerCAmelCase		)
									UpperCAmelCase__								=  config.num_labels
									UpperCAmelCase__								=  TFRegNetMainLayer(_lowerCAmelCase	,  name="regnet"		)
									# classification head
									UpperCAmelCase__								=  [
									    tf.keras.layers.Flatten(),
									    tf.keras.layers.Dense(config.num_labels	,  name="classifier.1"		) if config.num_labels > 0 else tf.identity,
									]





			@unpack_inputs
			@add_start_docstrings_to_model_forward(_lowerCAmelCase		)
			@add_code_sample_docstrings(
			    checkpoint=_IMAGE_CLASS_CHECKPOINT	,  output_type=_lowerCAmelCase	,  config_class=_CONFIG_FOR_DOC	,  expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT	,  )
			def       lowercase_				(self					:						List[Any]	,  __UpperCAmelCase					:						tf.Tensor = None	,  __UpperCAmelCase					:						tf.Tensor = None	,  __UpperCAmelCase					:						bool = None	,  __UpperCAmelCase					:						bool = None	,  __UpperCAmelCase					:						Dict=False	,  )						->    List[str]:
									"""simple docstring"""



									UpperCAmelCase__								=  (
									    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
									)
									UpperCAmelCase__								=  return_dict if return_dict is not None else self.config.use_return_dict

									UpperCAmelCase__								=  self.regnet(
									    _lowerCAmelCase	,  output_hidden_states=_lowerCAmelCase	,  return_dict=_lowerCAmelCase	,  training=_lowerCAmelCase		)

									UpperCAmelCase__								=  outputs.pooler_output if return_dict else outputs[1]

									UpperCAmelCase__								=  self.classifier[0](_lowerCAmelCase		)
									UpperCAmelCase__								=  self.classifier[1](_lowerCAmelCase		)

									UpperCAmelCase__								=  None if labels is None else self.hf_compute_loss(labels=_lowerCAmelCase	,  logits=_lowerCAmelCase		)

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

									return TFSequenceClassifierOutput(loss=_lowerCAmelCase	,  logits=_lowerCAmelCase	,  hidden_states=outputs.hidden_states		)




from __future__ import annotations

import unittest

from transformers import RoFormerConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
      import tensorflow as tf

      from transformers import (
          TFRoFormerForCausalLM,
          TFRoFormerForMaskedLM,
          TFRoFormerForMultipleChoice,
          TFRoFormerForQuestionAnswering,
          TFRoFormerForSequenceClassification,
          TFRoFormerForTokenClassification,
          TFRoFormerModel,
      )
      from transformers.models.roformer.modeling_tf_roformer import (
          TFRoFormerSelfAttention,
          TFRoFormerSinusoidalPositionalEmbedding,
      )





class 					A							:


   def __init__(self					:						int	,  __UpperCAmelCase					:						Tuple	,  __UpperCAmelCase					:						List[str]=1_3	,  __UpperCAmelCase					:						Optional[Any]=7	,  __UpperCAmelCase					:						Dict=True	,  __UpperCAmelCase					:						int=True	,  __UpperCAmelCase					:						Any=True	,  __UpperCAmelCase					:						int=True	,  __UpperCAmelCase					:						Dict=9_9	,  __UpperCAmelCase					:						Optional[int]=3_2	,  __UpperCAmelCase					:						Dict=2	,  __UpperCAmelCase					:						Any=4	,  __UpperCAmelCase					:						Tuple=3_7	,  __UpperCAmelCase					:						List[Any]="gelu"	,  __UpperCAmelCase					:						Dict=0.1	,  __UpperCAmelCase					:						int=0.1	,  __UpperCAmelCase					:						str=5_1_2	,  __UpperCAmelCase					:						Union[str, Any]=1_6	,  __UpperCAmelCase					:						Tuple=2	,  __UpperCAmelCase					:						Tuple=0.02	,  __UpperCAmelCase					:						Optional[int]=3	,  __UpperCAmelCase					:						str=4	,  __UpperCAmelCase					:						Any=None	,  )						->    Dict:
         """simple docstring"""



         UpperCAmelCase__								=  parent
         UpperCAmelCase__								=  1_3
         UpperCAmelCase__								=  7
         UpperCAmelCase__								=  True
         UpperCAmelCase__								=  True
         UpperCAmelCase__								=  True
         UpperCAmelCase__								=  True
         UpperCAmelCase__								=  9_9
         UpperCAmelCase__								=  3_2
         UpperCAmelCase__								=  2
         UpperCAmelCase__								=  4
         UpperCAmelCase__								=  3_7
         UpperCAmelCase__								=  "gelu"
         UpperCAmelCase__								=  0.1
         UpperCAmelCase__								=  0.1
         UpperCAmelCase__								=  5_1_2
         UpperCAmelCase__								=  1_6
         UpperCAmelCase__								=  2
         UpperCAmelCase__								=  0.02
         UpperCAmelCase__								=  3
         UpperCAmelCase__								=  4
         UpperCAmelCase__								=  None


   def       lowercase_				(self					:						Union[str, Any]		)						->    str:
         """simple docstring"""



         UpperCAmelCase__								=  ids_tensor([self.batch_size, self.seq_length]	,  self.vocab_size		)

         UpperCAmelCase__								=  None
         if self.use_input_mask:
               UpperCAmelCase__								=  random_attention_mask([self.batch_size, self.seq_length]		)

         UpperCAmelCase__								=  None
         if self.use_token_type_ids:
               UpperCAmelCase__								=  ids_tensor([self.batch_size, self.seq_length]	,  self.type_vocab_size		)

         UpperCAmelCase__								=  None
         UpperCAmelCase__								=  None
         UpperCAmelCase__								=  None
         if self.use_labels:
               UpperCAmelCase__								=  ids_tensor([self.batch_size]	,  self.type_sequence_label_size		)
               UpperCAmelCase__								=  ids_tensor([self.batch_size, self.seq_length]	,  self.num_labels		)
               UpperCAmelCase__								=  ids_tensor([self.batch_size]	,  self.num_choices		)

         UpperCAmelCase__								=  RoFormerConfig(
             vocab_size=self.vocab_size	,  hidden_size=self.hidden_size	,  num_hidden_layers=self.num_hidden_layers	,  num_attention_heads=self.num_attention_heads	,  intermediate_size=self.intermediate_size	,  hidden_act=self.hidden_act	,  hidden_dropout_prob=self.hidden_dropout_prob	,  attention_probs_dropout_prob=self.attention_probs_dropout_prob	,  max_position_embeddings=self.max_position_embeddings	,  type_vocab_size=self.type_vocab_size	,  initializer_range=self.initializer_range	,  return_dict=__UpperCAmelCase	,  )

         return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels


   def       lowercase_				(self					:						List[Any]	,  __UpperCAmelCase					:						str	,  __UpperCAmelCase					:						List[str]	,  __UpperCAmelCase					:						List[Any]	,  __UpperCAmelCase					:						List[str]	,  __UpperCAmelCase					:						List[str]	,  __UpperCAmelCase					:						Union[str, Any]	,  __UpperCAmelCase					:						Any		)						->    List[str]:
         """simple docstring"""



         UpperCAmelCase__								=  TFRoFormerModel(config=__UpperCAmelCase		)
         UpperCAmelCase__								=  {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}

         UpperCAmelCase__								=  [input_ids, input_mask]
         UpperCAmelCase__								=  model(__UpperCAmelCase		)

         UpperCAmelCase__								=  model(__UpperCAmelCase		)

         self.parent.assertEqual(result.last_hidden_state.shape	,  (self.batch_size, self.seq_length, self.hidden_size)		)


   def       lowercase_				(self					:						int	,  __UpperCAmelCase					:						Optional[int]	,  __UpperCAmelCase					:						Any	,  __UpperCAmelCase					:						Any	,  __UpperCAmelCase					:						List[str]	,  __UpperCAmelCase					:						Tuple	,  __UpperCAmelCase					:						Optional[int]	,  __UpperCAmelCase					:						List[Any]		)						->    Any:
         """simple docstring"""



         UpperCAmelCase__								=  True
         UpperCAmelCase__								=  TFRoFormerForCausalLM(config=__UpperCAmelCase		)
         UpperCAmelCase__								=  {
             "input_ids": input_ids,
             "attention_mask": input_mask,
             "token_type_ids": token_type_ids,
         }
         UpperCAmelCase__								=  model(__UpperCAmelCase		)["logits"]
         self.parent.assertListEqual(
             list(prediction_scores.numpy().shape		)	,  [self.batch_size, self.seq_length, self.vocab_size]		)


   def       lowercase_				(self					:						str	,  __UpperCAmelCase					:						Union[str, Any]	,  __UpperCAmelCase					:						int	,  __UpperCAmelCase					:						Optional[int]	,  __UpperCAmelCase					:						Any	,  __UpperCAmelCase					:						Optional[Any]	,  __UpperCAmelCase					:						Any	,  __UpperCAmelCase					:						int		)						->    Union[str, Any]:
         """simple docstring"""



         UpperCAmelCase__								=  TFRoFormerForMaskedLM(config=__UpperCAmelCase		)
         UpperCAmelCase__								=  {
             "input_ids": input_ids,
             "attention_mask": input_mask,
             "token_type_ids": token_type_ids,
         }
         UpperCAmelCase__								=  model(__UpperCAmelCase		)
         self.parent.assertEqual(result.logits.shape	,  (self.batch_size, self.seq_length, self.vocab_size)		)


   def       lowercase_				(self					:						Optional[Any]	,  __UpperCAmelCase					:						Optional[Any]	,  __UpperCAmelCase					:						Optional[int]	,  __UpperCAmelCase					:						Optional[Any]	,  __UpperCAmelCase					:						List[str]	,  __UpperCAmelCase					:						List[Any]	,  __UpperCAmelCase					:						Tuple	,  __UpperCAmelCase					:						str		)						->    Optional[Any]:
         """simple docstring"""



         UpperCAmelCase__								=  self.num_labels
         UpperCAmelCase__								=  TFRoFormerForSequenceClassification(config=__UpperCAmelCase		)
         UpperCAmelCase__								=  {
             "input_ids": input_ids,
             "attention_mask": input_mask,
             "token_type_ids": token_type_ids,
         }

         UpperCAmelCase__								=  model(__UpperCAmelCase		)
         self.parent.assertEqual(result.logits.shape	,  (self.batch_size, self.num_labels)		)


   def       lowercase_				(self					:						str	,  __UpperCAmelCase					:						Union[str, Any]	,  __UpperCAmelCase					:						int	,  __UpperCAmelCase					:						Any	,  __UpperCAmelCase					:						Tuple	,  __UpperCAmelCase					:						List[Any]	,  __UpperCAmelCase					:						Union[str, Any]	,  __UpperCAmelCase					:						Optional[Any]		)						->    Any:
         """simple docstring"""



         UpperCAmelCase__								=  self.num_choices
         UpperCAmelCase__								=  TFRoFormerForMultipleChoice(config=__UpperCAmelCase		)
         UpperCAmelCase__								=  tf.tile(tf.expand_dims(__UpperCAmelCase	,  1		)	,  (1, self.num_choices, 1)		)
         UpperCAmelCase__								=  tf.tile(tf.expand_dims(__UpperCAmelCase	,  1		)	,  (1, self.num_choices, 1)		)
         UpperCAmelCase__								=  tf.tile(tf.expand_dims(__UpperCAmelCase	,  1		)	,  (1, self.num_choices, 1)		)
         UpperCAmelCase__								=  {
             "input_ids": multiple_choice_inputs_ids,
             "attention_mask": multiple_choice_input_mask,
             "token_type_ids": multiple_choice_token_type_ids,
         }
         UpperCAmelCase__								=  model(__UpperCAmelCase		)
         self.parent.assertEqual(result.logits.shape	,  (self.batch_size, self.num_choices)		)


   def       lowercase_				(self					:						int	,  __UpperCAmelCase					:						Optional[Any]	,  __UpperCAmelCase					:						Tuple	,  __UpperCAmelCase					:						int	,  __UpperCAmelCase					:						int	,  __UpperCAmelCase					:						List[str]	,  __UpperCAmelCase					:						Union[str, Any]	,  __UpperCAmelCase					:						int		)						->    Union[str, Any]:
         """simple docstring"""



         UpperCAmelCase__								=  self.num_labels
         UpperCAmelCase__								=  TFRoFormerForTokenClassification(config=__UpperCAmelCase		)
         UpperCAmelCase__								=  {
             "input_ids": input_ids,
             "attention_mask": input_mask,
             "token_type_ids": token_type_ids,
         }
         UpperCAmelCase__								=  model(__UpperCAmelCase		)
         self.parent.assertEqual(result.logits.shape	,  (self.batch_size, self.seq_length, self.num_labels)		)


   def       lowercase_				(self					:						Dict	,  __UpperCAmelCase					:						Tuple	,  __UpperCAmelCase					:						Any	,  __UpperCAmelCase					:						Any	,  __UpperCAmelCase					:						Optional[Any]	,  __UpperCAmelCase					:						Optional[Any]	,  __UpperCAmelCase					:						Tuple	,  __UpperCAmelCase					:						Dict		)						->    List[str]:
         """simple docstring"""



         UpperCAmelCase__								=  TFRoFormerForQuestionAnswering(config=__UpperCAmelCase		)
         UpperCAmelCase__								=  {
             "input_ids": input_ids,
             "attention_mask": input_mask,
             "token_type_ids": token_type_ids,
         }

         UpperCAmelCase__								=  model(__UpperCAmelCase		)
         self.parent.assertEqual(result.start_logits.shape	,  (self.batch_size, self.seq_length)		)
         self.parent.assertEqual(result.end_logits.shape	,  (self.batch_size, self.seq_length)		)





   def       lowercase_				(self					:						Dict		)						->    Dict:
         """simple docstring"""



         UpperCAmelCase__								=  self.prepare_config_and_inputs()
         (
             (
             UpperCAmelCase__
         )		,				(
             UpperCAmelCase__
         )		,				(
             UpperCAmelCase__
         )		,				(
             UpperCAmelCase__
         )		,				(
             UpperCAmelCase__
         )		,				(
             UpperCAmelCase__
         )		,				(
             UpperCAmelCase__
         )		,				
         )								=  config_and_inputs
         UpperCAmelCase__								=  {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
         return config, inputs_dict





@require_tf
class 					A							(    UpperCAmelCase_		,    UpperCAmelCase_		,    unittest.TestCase					):
   __UpperCAmelCase			:			str              =      (
       (
           TFRoFormerModel,
           TFRoFormerForCausalLM,
           TFRoFormerForMaskedLM,
           TFRoFormerForQuestionAnswering,
           TFRoFormerForSequenceClassification,
           TFRoFormerForTokenClassification,
           TFRoFormerForMultipleChoice,
       )
       if is_tf_available()
       else ()
   )
   __UpperCAmelCase			:			List[str]              =      (
       {
           'feature-extraction': TFRoFormerModel,
           'fill-mask': TFRoFormerForMaskedLM,
           'question-answering': TFRoFormerForQuestionAnswering,
           'text-classification': TFRoFormerForSequenceClassification,
           'text-generation': TFRoFormerForCausalLM,
           'token-classification': TFRoFormerForTokenClassification,
           'zero-shot': TFRoFormerForSequenceClassification,
       }
       if is_tf_available()
       else {}
   )

   __UpperCAmelCase			:			Optional[int]              =      False
   __UpperCAmelCase			:			List[Any]              =      False


   def       lowercase_				(self					:						List[str]	,  __UpperCAmelCase					:						Union[str, Any]	,  __UpperCAmelCase					:						Dict	,  __UpperCAmelCase					:						Any	,  __UpperCAmelCase					:						Dict	,  __UpperCAmelCase					:						Any		)						->    int:
         """simple docstring"""



         if pipeline_test_casse_name == "TextGenerationPipelineTests":
               return True

         return False


   def       lowercase_				(self					:						Union[str, Any]		)						->    List[Any]:
         """simple docstring"""



         UpperCAmelCase__								=  TFRoFormerModelTester(self		)
         UpperCAmelCase__								=  ConfigTester(self	,  config_class=__UpperCAmelCase	,  hidden_size=3_7		)


   def       lowercase_				(self					:						Optional[int]		)						->    Any:
         """simple docstring"""



         self.config_tester.run_common_tests()


   def       lowercase_				(self					:						Union[str, Any]		)						->    List[Any]:
         """simple docstring"""



         UpperCAmelCase__								=  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*__UpperCAmelCase		)


   def       lowercase_				(self					:						Dict		)						->    Any:
         """simple docstring"""



         UpperCAmelCase__								=  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_masked_lm(*__UpperCAmelCase		)


   def       lowercase_				(self					:						Union[str, Any]		)						->    str:
         """simple docstring"""



         UpperCAmelCase__								=  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_lm_head(*__UpperCAmelCase		)


   def       lowercase_				(self					:						Optional[Any]		)						->    Dict:
         """simple docstring"""



         UpperCAmelCase__								=  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_multiple_choice(*__UpperCAmelCase		)


   def       lowercase_				(self					:						List[Any]		)						->    List[str]:
         """simple docstring"""



         UpperCAmelCase__								=  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_question_answering(*__UpperCAmelCase		)


   def       lowercase_				(self					:						int		)						->    List[Any]:
         """simple docstring"""



         UpperCAmelCase__								=  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_sequence_classification(*__UpperCAmelCase		)


   def       lowercase_				(self					:						List[str]		)						->    Tuple:
         """simple docstring"""



         UpperCAmelCase__								=  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_token_classification(*__UpperCAmelCase		)





   @slow
   def       lowercase_				(self					:						Any		)						->    Optional[Any]:
         """simple docstring"""



         UpperCAmelCase__								=  TFRoFormerModel.from_pretrained("junnyu/roformer_chinese_base"		)
         self.assertIsNotNone(__UpperCAmelCase		)





@require_tf
class 					A							(    unittest.TestCase					):


   @slow
   def       lowercase_				(self					:						List[Any]		)						->    Any:
         """simple docstring"""



         UpperCAmelCase__								=  TFRoFormerForMaskedLM.from_pretrained("junnyu/roformer_chinese_base"		)
         UpperCAmelCase__								=  tf.constant([[0, 1, 2, 3, 4, 5]]		)
         UpperCAmelCase__								=  model(__UpperCAmelCase		)[0]

         # TODO Replace vocab size
         UpperCAmelCase__								=  5_0_0_0_0

         UpperCAmelCase__								=  [1, 6, vocab_size]
         self.assertEqual(output.shape	,  __UpperCAmelCase		)

         print(output[:, :3, :3]		)

         # TODO Replace values below with what was printed above.
         UpperCAmelCase__								=  tf.constant(
             [
                 [
                     [-0.12053341, -1.0264901, 0.29221946],
                     [-1.5133783, 0.197433, 0.15190607],
                     [-5.0135403, -3.900256, -0.84038764],
                 ]
             ]		)
         tf.debugging.assert_near(output[:, :3, :3]	,  __UpperCAmelCase	,  atol=1E-4		)





@require_tf
class 					A							(    unittest.TestCase					):
   __UpperCAmelCase			:			Tuple              =      1E-4


   def       lowercase_				(self					:						List[str]		)						->    str:
         """simple docstring"""



         UpperCAmelCase__								=  tf.constant([[4, 1_0]]		)
         UpperCAmelCase__								=  TFRoFormerSinusoidalPositionalEmbedding(num_positions=6	,  embedding_dim=6		)

         UpperCAmelCase__								=  emba(input_ids.shape		)
         UpperCAmelCase__								=  tf.constant(
             [[0.0000, 0.0000, 0.0000, 1.0000, 1.0000, 1.0000], [0.8415, 0.0464, 0.0022, 0.5403, 0.9989, 1.0000]]		)

         tf.debugging.assert_near(__UpperCAmelCase	,  __UpperCAmelCase	,  atol=self.tolerance		)





   def       lowercase_				(self					:						List[Any]		)						->    int:
         """simple docstring"""



         UpperCAmelCase__								=  tf.constant(
             [
                 [0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
                 [0.8415, 0.8219, 0.8020, 0.7819, 0.7617],
                 [0.9093, 0.9364, 0.9581, 0.9749, 0.9870],
             ]		)
         UpperCAmelCase__								=  TFRoFormerSinusoidalPositionalEmbedding(num_positions=5_1_2	,  embedding_dim=5_1_2		)
         emba([2, 1_6, 5_1_2]		)
         UpperCAmelCase__								=  emba.weight[:3, :5]

         tf.debugging.assert_near(__UpperCAmelCase	,  __UpperCAmelCase	,  atol=self.tolerance		)





@require_tf
class 					A							(    unittest.TestCase					):
   __UpperCAmelCase			:			Any              =      1E-4


   def       lowercase_				(self					:						Tuple		)						->    int:
         """simple docstring"""



         UpperCAmelCase__								=  tf.reshape(tf.range(2 * 1_2 * 1_6 * 6_4	,  dtype=tf.floataa		)	,  shape=(2, 1_2, 1_6, 6_4)		) / 1_0_0

         UpperCAmelCase__								=  -tf.reshape(tf.range(2 * 1_2 * 1_6 * 6_4	,  dtype=tf.floataa		)	,  shape=(2, 1_2, 1_6, 6_4)		) / 1_0_0

         UpperCAmelCase__								=  TFRoFormerSinusoidalPositionalEmbedding(num_positions=3_2	,  embedding_dim=6_4		)
         UpperCAmelCase__								=  embed_positions([2, 1_6, 7_6_8]		)[None, None, :, :]

         UpperCAmelCase__		,				UpperCAmelCase__								=  TFRoFormerSelfAttention.apply_rotary_position_embeddings(
             __UpperCAmelCase	,  __UpperCAmelCase	,  __UpperCAmelCase		)

         UpperCAmelCase__								=  tf.constant(
             [
                 [0.0000, 0.0100, 0.0200, 0.0300, 0.0400, 0.0500, 0.0600, 0.0700],
                 [-0.2012, 0.8897, 0.0263, 0.9401, 0.2074, 0.9463, 0.3481, 0.9343],
                 [-1.7057, 0.6271, -1.2145, 1.3897, -0.6303, 1.7647, -0.1173, 1.8985],
                 [-2.1731, -1.6397, -2.7358, 0.2854, -2.1840, 1.7183, -1.3018, 2.4871],
                 [0.2717, -3.6173, -2.9206, -2.1988, -3.6638, 0.3858, -2.9155, 2.2980],
                 [3.9859, -2.1580, -0.7984, -4.4904, -4.1181, -2.0252, -4.4782, 1.1253],
             ]		)
         UpperCAmelCase__								=  tf.constant(
             [
                 [0.0000, -0.0100, -0.0200, -0.0300, -0.0400, -0.0500, -0.0600, -0.0700],
                 [0.2012, -0.8897, -0.0263, -0.9401, -0.2074, -0.9463, -0.3481, -0.9343],
                 [1.7057, -0.6271, 1.2145, -1.3897, 0.6303, -1.7647, 0.1173, -1.8985],
                 [2.1731, 1.6397, 2.7358, -0.2854, 2.1840, -1.7183, 1.3018, -2.4871],
                 [-0.2717, 3.6173, 2.9206, 2.1988, 3.6638, -0.3858, 2.9155, -2.2980],
                 [-3.9859, 2.1580, 0.7984, 4.4904, 4.1181, 2.0252, 4.4782, -1.1253],
             ]		)

         tf.debugging.assert_near(query_layer[0, 0, :6, :8]	,  __UpperCAmelCase	,  atol=self.tolerance		)
         tf.debugging.assert_near(key_layer[0, 0, :6, :8]	,  __UpperCAmelCase	,  atol=self.tolerance		)



