
"""simple docstring"""

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


SCREAMING_SNAKE_CASE_            =    logging.get_logger(__name__)

# General docstring
SCREAMING_SNAKE_CASE_            =    '''RegNetConfig'''

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

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

SCREAMING_SNAKE_CASE_            =    [
    '''facebook/regnet-y-040''',
    # See all regnet models at https://huggingface.co/models?filter=regnet
]
class   _UpperCAmelCase				(		tf.keras.layers.Layer				):




      def __init__(				self ,							lowercase_ ,							lowercase_ = 3 ,							lowercase_ = 1 ,							lowercase_ = 1 ,							lowercase_ = "relu" ,							**lowercase_ ,							)    ->       Any:
          super().__init__(**lowercase_				)
          # 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=lowercase_ ,							kernel_size=lowercase_ ,							strides=lowercase_ ,							padding='VALID' ,							groups=lowercase_ ,							use_bias=lowercase_ ,							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      a_   (				self ,							lowercase_				)    ->       str:
          UpperCAmelCase          =       self.convolution(self.padding(lowercase_				)				)
          UpperCAmelCase          =       self.normalization(lowercase_				)
          UpperCAmelCase          =       self.activation(lowercase_				)
          return hidden_state

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




      def __init__(				self ,							lowercase_ ,							**lowercase_				)    ->       Optional[int]:
          super().__init__(**lowercase_				)
          UpperCAmelCase          =       config.num_channels
          UpperCAmelCase          =       TFRegNetConvLayer(
              out_channels=config.embedding_size ,							kernel_size=3 ,							stride=2 ,							activation=config.hidden_act ,							name='embedder' ,							)




      def      a_   (				self ,							lowercase_				)    ->       List[str]:
          UpperCAmelCase          =       shape_list(lowercase_				)[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(lowercase_ ,							perm=(0, 2, 3, 1)				)
          UpperCAmelCase          =       self.embedder(lowercase_				)
          return hidden_state

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




      def __init__(				self ,							lowercase_ ,							lowercase_ = 2 ,							**lowercase_				)    ->       str:
          super().__init__(**lowercase_				)
          UpperCAmelCase          =       tf.keras.layers.ConvaD(
              filters=lowercase_ ,							kernel_size=1 ,							strides=lowercase_ ,							use_bias=lowercase_ ,							name='convolution'				)
          UpperCAmelCase          =       tf.keras.layers.BatchNormalization(epsilon=1E-5 ,							momentum=0.9 ,							name='normalization'				)




      def      a_   (				self ,							lowercase_ ,							lowercase_ = False				)    ->       List[Any]:
          return self.normalization(self.convolution(lowercase_				) ,							training=lowercase_				)

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




      def __init__(				self ,							lowercase_ ,							lowercase_ ,							**lowercase_				)    ->       Dict:
          super().__init__(**lowercase_				)
          UpperCAmelCase          =       tf.keras.layers.GlobalAveragePoolingaD(keepdims=lowercase_ ,							name='pooler'				)
          UpperCAmelCase          =       [
              tf.keras.layers.ConvaD(filters=lowercase_ ,							kernel_size=1 ,							activation='relu' ,							name='attention.0'				),
              tf.keras.layers.ConvaD(filters=lowercase_ ,							kernel_size=1 ,							activation='sigmoid' ,							name='attention.2'				),
          ]




      def      a_   (				self ,							lowercase_				)    ->       List[Any]:
          UpperCAmelCase          =       self.pooler(lowercase_				)
          for layer_module in self.attention:
              UpperCAmelCase          =       layer_module(lowercase_				)
          UpperCAmelCase          =       hidden_state * pooled
          return hidden_state

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




      def __init__(				self ,							lowercase_ ,							lowercase_ ,							lowercase_ ,							lowercase_ = 1 ,							**lowercase_				)    ->       List[str]:
          super().__init__(**lowercase_				)
          UpperCAmelCase          =       in_channels != out_channels or stride != 1
          UpperCAmelCase          =       max(1 ,							out_channels // config.groups_width				)
          UpperCAmelCase          =       (
              TFRegNetShortCut(lowercase_ ,							stride=lowercase_ ,							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(lowercase_ ,							kernel_size=1 ,							activation=config.hidden_act ,							name='layer.0'				),
              TFRegNetConvLayer(
                  lowercase_ ,							stride=lowercase_ ,							groups=lowercase_ ,							activation=config.hidden_act ,							name='layer.1'				),
              TFRegNetConvLayer(lowercase_ ,							kernel_size=1 ,							activation=lowercase_ ,							name='layer.2'				),
          ]
          UpperCAmelCase          =       ACTaFN[config.hidden_act]




      def      a_   (				self ,							lowercase_				)    ->       List[Any]:
          UpperCAmelCase          =       hidden_state
          for layer_module in self.layers:
              UpperCAmelCase          =       layer_module(lowercase_				)
          UpperCAmelCase          =       self.shortcut(lowercase_				)
          hidden_state += residual
          UpperCAmelCase          =       self.activation(lowercase_				)
          return hidden_state

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




      def __init__(				self ,							lowercase_ ,							lowercase_ ,							lowercase_ ,							lowercase_ = 1 ,							**lowercase_				)    ->       Union[str, Any]:
          super().__init__(**lowercase_				)
          UpperCAmelCase          =       in_channels != out_channels or stride != 1
          UpperCAmelCase          =       max(1 ,							out_channels // config.groups_width				)
          UpperCAmelCase          =       (
              TFRegNetShortCut(lowercase_ ,							stride=lowercase_ ,							name='shortcut'				)
              if should_apply_shortcut
              else tf.keras.layers.Activation('linear' ,							name='shortcut'				)
          )
          UpperCAmelCase          =       [
              TFRegNetConvLayer(lowercase_ ,							kernel_size=1 ,							activation=config.hidden_act ,							name='layer.0'				),
              TFRegNetConvLayer(
                  lowercase_ ,							stride=lowercase_ ,							groups=lowercase_ ,							activation=config.hidden_act ,							name='layer.1'				),
              TFRegNetSELayer(lowercase_ ,							reduced_channels=int(round(in_channels / 4				)				) ,							name='layer.2'				),
              TFRegNetConvLayer(lowercase_ ,							kernel_size=1 ,							activation=lowercase_ ,							name='layer.3'				),
          ]
          UpperCAmelCase          =       ACTaFN[config.hidden_act]




      def      a_   (				self ,							lowercase_				)    ->       Dict:
          UpperCAmelCase          =       hidden_state
          for layer_module in self.layers:
              UpperCAmelCase          =       layer_module(lowercase_				)
          UpperCAmelCase          =       self.shortcut(lowercase_				)
          hidden_state += residual
          UpperCAmelCase          =       self.activation(lowercase_				)
          return hidden_state

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




      def __init__(				self ,							lowercase_ ,							lowercase_ ,							lowercase_ ,							lowercase_ = 2 ,							lowercase_ = 2 ,							**lowercase_				)    ->       Optional[Any]:
          super().__init__(**lowercase_				)

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




      def      a_   (				self ,							lowercase_				)    ->       Union[str, Any]:





          for layer_module in self.layers:
              UpperCAmelCase          =       layer_module(lowercase_				)
          return hidden_state

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




      def __init__(				self ,							lowercase_ ,							**lowercase_				)    ->       Optional[int]:
          super().__init__(**lowercase_				)
          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(
                  lowercase_ ,							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(lowercase_ ,							config.depths[1:]				)				):
              self.stages.append(TFRegNetStage(lowercase_ ,							lowercase_ ,							lowercase_ ,							depth=lowercase_ ,							name=F"stages.{i+1}"				)				)




      def      a_   (				self ,							lowercase_ ,							lowercase_ = False ,							lowercase_ = True				)    ->       List[str]:
          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(lowercase_				)

          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=lowercase_ ,							hidden_states=lowercase_				)

@keras_serializable
class   _UpperCAmelCase				(		tf.keras.layers.Layer				):
      __SCREAMING_SNAKE_CASE :				Optional[int]				   =						RegNetConfig




      def __init__(				self ,							lowercase_ ,							**lowercase_				)    ->       str:
          super().__init__(**lowercase_				)
          UpperCAmelCase          =       config
          UpperCAmelCase          =       TFRegNetEmbeddings(lowercase_ ,							name='embedder'				)
          UpperCAmelCase          =       TFRegNetEncoder(lowercase_ ,							name='encoder'				)
          UpperCAmelCase          =       tf.keras.layers.GlobalAveragePoolingaD(keepdims=lowercase_ ,							name='pooler'				)




      @unpack_inputs
      def      a_   (				self ,							lowercase_ ,							lowercase_ = None ,							lowercase_ = None ,							lowercase_ = False ,							)    ->       Optional[Any]:
          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(lowercase_ ,							training=lowercase_				)

          UpperCAmelCase          =       self.encoder(
              lowercase_ ,							output_hidden_states=lowercase_ ,							return_dict=lowercase_ ,							training=lowercase_				)

          UpperCAmelCase          =       encoder_outputs[0]
          UpperCAmelCase          =       self.pooler(lowercase_				)

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

          # Change the other hidden state outputs to NCHW as well
          if output_hidden_states:
              UpperCAmelCase          =       tuple([tf.transpose(lowercase_ ,							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=lowercase_ ,							pooler_output=lowercase_ ,							hidden_states=hidden_states if output_hidden_states else encoder_outputs.hidden_states ,							)

class   _UpperCAmelCase				(		SCREAMING_SNAKE_CASE_				):
      __SCREAMING_SNAKE_CASE :				Dict				   =						RegNetConfig
      __SCREAMING_SNAKE_CASE :				List[str]				   =						"regnet"
      __SCREAMING_SNAKE_CASE :				List[str]				   =						"pixel_values"




      @property
      def      a_   (				self				)    ->       int:
          return {"pixel_values": tf.TensorSpec(shape=(None, self.config.num_channels, 2_2_4, 2_2_4) ,							dtype=tf.floataa				)}

SCREAMING_SNAKE_CASE_            =    r'''
    Parameters:
    This model is a Tensorflow
    [tf.keras.layers.Layer](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Layer) sub-class. Use it as a
    regular Tensorflow Module and refer to the Tensorflow documentation for all matter related to general usage and
    behavior.
        config ([`RegNetConfig`]): 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 [`~TFPreTrainedModel.from_pretrained`] method to load the model weights.
'''

SCREAMING_SNAKE_CASE_            =    r'''
    Args:
        pixel_values (`tf.Tensor` of shape `(batch_size, num_channels, height, width)`):
            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See
            [`ConveNextImageProcessor.__call__`] for details.
        output_hidden_states (`bool`, *optional*):
            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
            more detail.
        return_dict (`bool`, *optional*):
            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
'''
@add_start_docstrings(
    "The bare RegNet model outputting raw features without any specific head on top."  ,      SCREAMING_SNAKE_CASE_  ,      )
class   _UpperCAmelCase				(		SCREAMING_SNAKE_CASE_				):




      def __init__(				self ,							lowercase_ ,							*lowercase_ ,							**lowercase_				)    ->       str:
          super().__init__(lowercase_ ,							*lowercase_ ,							**lowercase_				)
          UpperCAmelCase          =       TFRegNetMainLayer(lowercase_ ,							name='regnet'				)




      @unpack_inputs
      @add_start_docstrings_to_model_forward(lowercase_				)
      @add_code_sample_docstrings(
          checkpoint=_CHECKPOINT_FOR_DOC ,							output_type=lowercase_ ,							config_class=_CONFIG_FOR_DOC ,							modality='vision' ,							expected_output=_EXPECTED_OUTPUT_SHAPE ,							)
      def      a_   (				self ,							lowercase_ ,							lowercase_ = None ,							lowercase_ = None ,							lowercase_=False ,							)    ->       Union[str, Any]:
          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=lowercase_ ,							output_hidden_states=lowercase_ ,							return_dict=lowercase_ ,							training=lowercase_ ,							)
          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    "  ,      SCREAMING_SNAKE_CASE_  ,      )
class   _UpperCAmelCase				(		SCREAMING_SNAKE_CASE_  ,      SCREAMING_SNAKE_CASE_				):




      def __init__(				self ,							lowercase_ ,							*lowercase_ ,							**lowercase_				)    ->       List[Any]:
          super().__init__(lowercase_ ,							*lowercase_ ,							**lowercase_				)
          UpperCAmelCase          =       config.num_labels
          UpperCAmelCase          =       TFRegNetMainLayer(lowercase_ ,							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(lowercase_				)
      @add_code_sample_docstrings(
          checkpoint=_IMAGE_CLASS_CHECKPOINT ,							output_type=lowercase_ ,							config_class=_CONFIG_FOR_DOC ,							expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT ,							)
      def      a_   (				self ,							lowercase_ = None ,							lowercase_ = None ,							lowercase_ = None ,							lowercase_ = None ,							lowercase_=False ,							)    ->       Tuple:
          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(
              lowercase_ ,							output_hidden_states=lowercase_ ,							return_dict=lowercase_ ,							training=lowercase_				)

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

          UpperCAmelCase          =       self.classifier[0](lowercase_				)
          UpperCAmelCase          =       self.classifier[1](lowercase_				)

          UpperCAmelCase          =       None if labels is None else self.hf_compute_loss(labels=lowercase_ ,							logits=lowercase_				)

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

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







'''simple docstring'''


import argparse

import torch

from transformers import GPTaLMHeadModel, RobertaForMaskedLM


if __name__ == "__main__":
      _UpperCamelCase 							=  argparse.ArgumentParser(
          description=(
              """Extraction some layers of the full RobertaForMaskedLM or GPT2LMHeadModel for Transfer Learned"""
              """ Distillation"""
          )
      )
      parser.add_argument("""--model_type""", default="""roberta""", choices=["""roberta""", """gpt2"""])
      parser.add_argument("""--model_name""", default="""roberta-large""", type=str)
      parser.add_argument("""--dump_checkpoint""", default="""serialization_dir/tf_roberta_048131723.pth""", type=str)
      parser.add_argument("""--vocab_transform""", action="""store_true""")
      _UpperCamelCase 							=  parser.parse_args()

      if args.model_type == "roberta":
            _UpperCamelCase 							=  RobertaForMaskedLM.from_pretrained(args.model_name)
            _UpperCamelCase 							=  """roberta"""
      elif args.model_type == "gpt2":
            _UpperCamelCase 							=  GPTaLMHeadModel.from_pretrained(args.model_name)
            _UpperCamelCase 							=  """transformer"""

      _UpperCamelCase 							=  model.state_dict()
      _UpperCamelCase 							=  {}

      # Embeddings #
      if args.model_type == "gpt2":
            for param_name in ["wte.weight", "wpe.weight"]:
                  _UpperCamelCase 							=  state_dict[F"""{prefix}.{param_name}"""]
      else:
            for w in ["word_embeddings", "position_embeddings", "token_type_embeddings"]:
                  _UpperCamelCase 							=  F"""{prefix}.embeddings.{w}.weight"""
                  _UpperCamelCase 							=  state_dict[param_name]
            for w in ["weight", "bias"]:
                  _UpperCamelCase 							=  F"""{prefix}.embeddings.LayerNorm.{w}"""
                  _UpperCamelCase 							=  state_dict[param_name]

    # Transformer Blocks #
      _UpperCamelCase 							=  0
      for teacher_idx in [0, 2, 4, 7, 9, 11]:
            if args.model_type == "gpt2":
                  for layer in ["ln_1", "attn.c_attn", "attn.c_proj", "ln_2", "mlp.c_fc", "mlp.c_proj"]:
                        for w in ["weight", "bias"]:
                              _UpperCamelCase 							=  state_dict[
                                  F"""{prefix}.h.{teacher_idx}.{layer}.{w}"""
                              ]
                  _UpperCamelCase 							=  state_dict[F"""{prefix}.h.{teacher_idx}.attn.bias"""]
            else:
                  for layer in [
                      "attention.self.query",
                      "attention.self.key",
                      "attention.self.value",
                      "attention.output.dense",
                      "attention.output.LayerNorm",
                      "intermediate.dense",
                      "output.dense",
                      "output.LayerNorm",
                  ]:
                        for w in ["weight", "bias"]:
                              _UpperCamelCase 							=  state_dict[
                                  F"""{prefix}.encoder.layer.{teacher_idx}.{layer}.{w}"""
                              ]
            std_idx += 1

      # Language Modeling Head ###s
      if args.model_type == "roberta":
            for layer in ["lm_head.decoder.weight", "lm_head.bias"]:
                  _UpperCamelCase 							=  state_dict[F"""{layer}"""]
            if args.vocab_transform:
                  for w in ["weight", "bias"]:
                        _UpperCamelCase 							=  state_dict[F"""lm_head.dense.{w}"""]
                        _UpperCamelCase 							=  state_dict[F"""lm_head.layer_norm.{w}"""]
      elif args.model_type == "gpt2":
            for w in ["weight", "bias"]:
                  _UpperCamelCase 							=  state_dict[F"""{prefix}.ln_f.{w}"""]
            _UpperCamelCase 							=  state_dict["""lm_head.weight"""]

      print(F"""N layers selected for distillation: {std_idx}""")
      print(F"""Number of params transferred for distillation: {len(compressed_sd.keys())}""")

      print(F"""Save transferred checkpoint to {args.dump_checkpoint}.""")
      torch.save(compressed_sd, args.dump_checkpoint)

