'''simple docstring'''


from __future__ import annotations

import unittest

from transformers import FunnelConfig, is_tf_available
from transformers.testing_utils import require_tf

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 (
           TFFunnelBaseModel,
           TFFunnelForMaskedLM,
           TFFunnelForMultipleChoice,
           TFFunnelForPreTraining,
           TFFunnelForQuestionAnswering,
           TFFunnelForSequenceClassification,
           TFFunnelForTokenClassification,
           TFFunnelModel,
       )





class       _a    :


  """simple docstring"""

  def __init__(		self      :       Any   ,				lowercase_      :       Dict   ,				lowercase_      :       int=13   ,				lowercase_      :       Optional[int]=7   ,				lowercase_      :       Tuple=True   ,				lowercase_      :       List[str]=True   ,				lowercase_      :       int=True   ,				lowercase_      :       str=True   ,				lowercase_      :       List[Any]=99   ,				lowercase_      :       List[str]=[1, 1, 2]   ,				lowercase_      :       str=1   ,				lowercase_      :       Union[str, Any]=32   ,				lowercase_      :       Tuple=4   ,				lowercase_      :       Optional[Any]=8   ,				lowercase_      :       Dict=37   ,				lowercase_      :       str="gelu_new"   ,				lowercase_      :       List[str]=0.1   ,				lowercase_      :       Tuple=0.1   ,				lowercase_      :       int=0.0   ,				lowercase_      :       Union[str, Any]=512   ,				lowercase_      :       Tuple=3   ,				lowercase_      :       str=0.0_2   ,				lowercase_      :       Optional[Any]=3   ,				lowercase_      :       str=4   ,				lowercase_      :       Optional[int]=None   ,				lowercase_      :       Optional[Any]=False   ,				):







    '''simple docstring'''


    lowercase_     =       parent
    lowercase_     =       batch_size
    lowercase_     =       seq_length
    lowercase_     =       is_training
    lowercase_     =       use_input_mask
    lowercase_     =       use_token_type_ids
    lowercase_     =       use_labels
    lowercase_     =       vocab_size
    lowercase_     =       block_sizes
    lowercase_     =       num_decoder_layers
    lowercase_     =       d_model
    lowercase_     =       n_head
    lowercase_     =       d_head
    lowercase_     =       d_inner
    lowercase_     =       hidden_act
    lowercase_     =       hidden_dropout
    lowercase_     =       attention_dropout
    lowercase_     =       activation_dropout
    lowercase_     =       max_position_embeddings
    lowercase_     =       type_vocab_size
    lowercase_     =       2
    lowercase_     =       num_labels
    lowercase_     =       num_choices
    lowercase_     =       scope
    lowercase_     =       initializer_std

    # Used in the tests to check the size of the first attention layer
    lowercase_     =       n_head
    # Used in the tests to check the size of the first hidden state
    lowercase_     =       self.d_model
    # Used in the tests to check the number of output hidden states/attentions
    lowercase_     =       sum(self.block_sizes      ) + (0 if base else self.num_decoder_layers)
    # FunnelModel adds two hidden layers: input embeddings and the sum of the upsampled encoder hidden state with
    # the last hidden state of the first block (which is the first hidden state of the decoder).
    if not base:
      lowercase_     =       self.num_hidden_layers + 2




  def 		lowerCamelCase__  (		self      :       str      ):







    '''simple docstring'''


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

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

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

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

    lowercase_     =       FunnelConfig(
        vocab_size=self.vocab_size   ,				block_sizes=self.block_sizes   ,				num_decoder_layers=self.num_decoder_layers   ,				d_model=self.d_model   ,				n_head=self.n_head   ,				d_head=self.d_head   ,				d_inner=self.d_inner   ,				hidden_act=self.hidden_act   ,				hidden_dropout=self.hidden_dropout   ,				attention_dropout=self.attention_dropout   ,				activation_dropout=self.activation_dropout   ,				max_position_embeddings=self.max_position_embeddings   ,				type_vocab_size=self.type_vocab_size   ,				initializer_std=self.initializer_std   ,				)

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




  def 		lowerCamelCase__  (		self      :       Union[str, Any]   ,				lowercase_      :       List[str]   ,				lowercase_      :       List[str]   ,				lowercase_      :       List[Any]   ,				lowercase_      :       Optional[Any]   ,				lowercase_      :       int   ,				lowercase_      :       Union[str, Any]   ,				lowercase_      :       str   ,				):







    '''simple docstring'''


    lowercase_     =       TFFunnelModel(config=lowercase_      )
    lowercase_     =       {"""input_ids""": input_ids, """attention_mask""": input_mask, """token_type_ids""": token_type_ids}
    lowercase_     =       model(lowercase_      )

    lowercase_     =       [input_ids, input_mask]
    lowercase_     =       model(lowercase_      )

    lowercase_     =       model(lowercase_      )
    self.parent.assertEqual(result.last_hidden_state.shape   ,				(self.batch_size, self.seq_length, self.d_model)      )

    lowercase_     =       False
    lowercase_     =       TFFunnelModel(config=lowercase_      )
    lowercase_     =       model(lowercase_      )
    self.parent.assertEqual(result.last_hidden_state.shape   ,				(self.batch_size, self.seq_length, self.d_model)      )

    lowercase_     =       False
    lowercase_     =       TFFunnelModel(config=lowercase_      )
    lowercase_     =       model(lowercase_      )
    self.parent.assertEqual(result.last_hidden_state.shape   ,				(self.batch_size, self.seq_length, self.d_model)      )




  def 		lowerCamelCase__  (		self      :       List[Any]   ,				lowercase_      :       Any   ,				lowercase_      :       Union[str, Any]   ,				lowercase_      :       Optional[int]   ,				lowercase_      :       Optional[int]   ,				lowercase_      :       Any   ,				lowercase_      :       int   ,				lowercase_      :       Any   ,				):







    '''simple docstring'''


    lowercase_     =       TFFunnelBaseModel(config=lowercase_      )
    lowercase_     =       {"""input_ids""": input_ids, """attention_mask""": input_mask, """token_type_ids""": token_type_ids}
    lowercase_     =       model(lowercase_      )

    lowercase_     =       [input_ids, input_mask]
    lowercase_     =       model(lowercase_      )

    lowercase_     =       model(lowercase_      )
    self.parent.assertEqual(result.last_hidden_state.shape   ,				(self.batch_size, 2, self.d_model)      )

    lowercase_     =       False
    lowercase_     =       TFFunnelBaseModel(config=lowercase_      )
    lowercase_     =       model(lowercase_      )
    self.parent.assertEqual(result.last_hidden_state.shape   ,				(self.batch_size, 3, self.d_model)      )

    lowercase_     =       False
    lowercase_     =       TFFunnelBaseModel(config=lowercase_      )
    lowercase_     =       model(lowercase_      )
    self.parent.assertEqual(result.last_hidden_state.shape   ,				(self.batch_size, 2, self.d_model)      )




  def 		lowerCamelCase__  (		self      :       Dict   ,				lowercase_      :       Any   ,				lowercase_      :       Optional[Any]   ,				lowercase_      :       Union[str, Any]   ,				lowercase_      :       List[str]   ,				lowercase_      :       int   ,				lowercase_      :       Any   ,				lowercase_      :       Dict   ,				):







    '''simple docstring'''


    lowercase_     =       TFFunnelForPreTraining(config=lowercase_      )
    lowercase_     =       {"""input_ids""": input_ids, """attention_mask""": input_mask, """token_type_ids""": token_type_ids}
    lowercase_     =       model(lowercase_      )
    self.parent.assertEqual(result.logits.shape   ,				(self.batch_size, self.seq_length)      )




  def 		lowerCamelCase__  (		self      :       Union[str, Any]   ,				lowercase_      :       Dict   ,				lowercase_      :       Any   ,				lowercase_      :       Union[str, Any]   ,				lowercase_      :       List[Any]   ,				lowercase_      :       Dict   ,				lowercase_      :       List[Any]   ,				lowercase_      :       Optional[int]   ,				):







    '''simple docstring'''


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




  def 		lowerCamelCase__  (		self      :       int   ,				lowercase_      :       List[str]   ,				lowercase_      :       str   ,				lowercase_      :       str   ,				lowercase_      :       Dict   ,				lowercase_      :       Dict   ,				lowercase_      :       int   ,				lowercase_      :       int   ,				):







    '''simple docstring'''


    lowercase_     =       self.num_labels
    lowercase_     =       TFFunnelForSequenceClassification(config=lowercase_      )
    lowercase_     =       {"""input_ids""": input_ids, """attention_mask""": input_mask, """token_type_ids""": token_type_ids}
    lowercase_     =       model(lowercase_      )
    self.parent.assertEqual(result.logits.shape   ,				(self.batch_size, self.num_labels)      )




  def 		lowerCamelCase__  (		self      :       Any   ,				lowercase_      :       int   ,				lowercase_      :       Tuple   ,				lowercase_      :       Optional[int]   ,				lowercase_      :       Union[str, Any]   ,				lowercase_      :       Any   ,				lowercase_      :       Dict   ,				lowercase_      :       Optional[Any]   ,				):







    '''simple docstring'''


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




  def 		lowerCamelCase__  (		self      :       int   ,				lowercase_      :       Optional[int]   ,				lowercase_      :       Union[str, Any]   ,				lowercase_      :       List[str]   ,				lowercase_      :       Union[str, Any]   ,				lowercase_      :       List[Any]   ,				lowercase_      :       Optional[Any]   ,				lowercase_      :       Dict   ,				):







    '''simple docstring'''


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




  def 		lowerCamelCase__  (		self      :       Optional[Any]   ,				lowercase_      :       Union[str, Any]   ,				lowercase_      :       Tuple   ,				lowercase_      :       Optional[Any]   ,				lowercase_      :       int   ,				lowercase_      :       List[str]   ,				lowercase_      :       List[str]   ,				lowercase_      :       List[Any]   ,				):







    '''simple docstring'''


    lowercase_     =       TFFunnelForQuestionAnswering(config=lowercase_      )
    lowercase_     =       {"""input_ids""": input_ids, """attention_mask""": input_mask, """token_type_ids""": token_type_ids}
    lowercase_     =       model(lowercase_      )
    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 		lowerCamelCase__  (		self      :       Union[str, Any]      ):







    '''simple docstring'''


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






@require_tf
class       _a    (		__a ,				__a ,				unittest.TestCase      ):


  """simple docstring"""

  A_    							=						(
      (
          TFFunnelModel,
          TFFunnelForMaskedLM,
          TFFunnelForPreTraining,
          TFFunnelForQuestionAnswering,
          TFFunnelForTokenClassification,
      )
      if is_tf_available()
      else ()
  )
  A_    							=						(
      {
          '''feature-extraction''': (TFFunnelBaseModel, TFFunnelModel),
          '''fill-mask''': TFFunnelForMaskedLM,
          '''question-answering''': TFFunnelForQuestionAnswering,
          '''text-classification''': TFFunnelForSequenceClassification,
          '''token-classification''': TFFunnelForTokenClassification,
          '''zero-shot''': TFFunnelForSequenceClassification,
      }
      if is_tf_available()
      else {}
  )
  A_    							=						False
  A_    							=						False




  def 		lowerCamelCase__  (		self      :       str      ):







    '''simple docstring'''


    lowercase_     =       TFFunnelModelTester(self      )
    lowercase_     =       ConfigTester(self   ,				config_class=lowercase_      )




  def 		lowerCamelCase__  (		self      :       List[Any]      ):







    '''simple docstring'''


    self.config_tester.run_common_tests()




  def 		lowerCamelCase__  (		self      :       int      ):







    '''simple docstring'''


    lowercase_     =       self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_model(*lowercase_      )




  def 		lowerCamelCase__  (		self      :       Optional[Any]      ):







    '''simple docstring'''


    lowercase_     =       self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_pretraining(*lowercase_      )




  def 		lowerCamelCase__  (		self      :       str      ):







    '''simple docstring'''


    lowercase_     =       self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_masked_lm(*lowercase_      )




  def 		lowerCamelCase__  (		self      :       List[str]      ):







    '''simple docstring'''


    lowercase_     =       self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_token_classification(*lowercase_      )





  def 		lowerCamelCase__  (		self      :       Dict      ):







    '''simple docstring'''


    lowercase_     =       self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_question_answering(*lowercase_      )






@require_tf
class       _a    (		__a ,				unittest.TestCase      ):


  """simple docstring"""

  A_    							=						(
      (TFFunnelBaseModel, TFFunnelForMultipleChoice, TFFunnelForSequenceClassification) if is_tf_available() else ()
  )
  A_    							=						False
  A_    							=						False




  def 		lowerCamelCase__  (		self      :       Optional[int]      ):







    '''simple docstring'''


    lowercase_     =       TFFunnelModelTester(self   ,				base=lowercase_      )
    lowercase_     =       ConfigTester(self   ,				config_class=lowercase_      )




  def 		lowerCamelCase__  (		self      :       Dict      ):







    '''simple docstring'''


    self.config_tester.run_common_tests()




  def 		lowerCamelCase__  (		self      :       Any      ):







    '''simple docstring'''


    lowercase_     =       self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_base_model(*lowercase_      )




  def 		lowerCamelCase__  (		self      :       Dict      ):







    '''simple docstring'''


    lowercase_     =       self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_sequence_classification(*lowercase_      )





  def 		lowerCamelCase__  (		self      :       int      ):







    '''simple docstring'''


    lowercase_     =       self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_multiple_choice(*lowercase_      )


'''simple docstring'''


import argparse

import torch
from torch import nn

from transformers import SpeechaTextConfig, SpeechaTextForConditionalGeneration

def   A_      (       SCREAMING_SNAKE_CASE_						)			->List[Any]:
  lowercase_     =       [
      """encoder.version""",
      """decoder.version""",
      """model.encoder.version""",
      """model.decoder.version""",
      """decoder.output_projection.weight""",
      """_float_tensor""",
      """encoder.embed_positions._float_tensor""",
      """decoder.embed_positions._float_tensor""",
  ]
  for k in ignore_keys:
    state_dict.pop(SCREAMING_SNAKE_CASE_ ,	SCREAMING_SNAKE_CASE_						)

def   A_      (       SCREAMING_SNAKE_CASE_						)			->Optional[int]:
  lowercase_     =       list(s_dict.keys()						)
  for key in keys:
    if "transformer_layers" in key:
      lowercase_     =       s_dict.pop(SCREAMING_SNAKE_CASE_						)
    elif "subsample" in key:
      lowercase_     =       s_dict.pop(SCREAMING_SNAKE_CASE_						)

def   A_      (       SCREAMING_SNAKE_CASE_						)			->Dict:
  lowercase_     ,	lowercase_     =       emb.weight.shape
  lowercase_     =       nn.Linear(SCREAMING_SNAKE_CASE_ ,	SCREAMING_SNAKE_CASE_ ,	bias=SCREAMING_SNAKE_CASE_						)
  lowercase_     =       emb.weight.data
  return lin_layer

def   A_      (       SCREAMING_SNAKE_CASE_ ,	SCREAMING_SNAKE_CASE_						)			->Any:
  lowercase_     =       torch.load(SCREAMING_SNAKE_CASE_ ,	map_location="""cpu"""						)
  lowercase_     =       mam_aaa["""args"""]
  lowercase_     =       mam_aaa["""model"""]
  lowercase_     =       state_dict["""decoder.output_projection.weight"""]

  remove_ignore_keys_(SCREAMING_SNAKE_CASE_						)
  rename_keys(SCREAMING_SNAKE_CASE_						)

  lowercase_     =       state_dict["""decoder.embed_tokens.weight"""].shape[0]

  lowercase_     =       args.share_decoder_input_output_embed

  lowercase_     =       [int(SCREAMING_SNAKE_CASE_						) for i in args.conv_kernel_sizes.split(""","""						)]
  lowercase_     =       SpeechaTextConfig(
      vocab_size=SCREAMING_SNAKE_CASE_ ,	max_source_positions=args.max_source_positions ,	max_target_positions=args.max_target_positions ,	encoder_layers=args.encoder_layers ,	decoder_layers=args.decoder_layers ,	encoder_attention_heads=args.encoder_attention_heads ,	decoder_attention_heads=args.decoder_attention_heads ,	encoder_ffn_dim=args.encoder_ffn_embed_dim ,	decoder_ffn_dim=args.decoder_ffn_embed_dim ,	d_model=args.encoder_embed_dim ,	dropout=args.dropout ,	attention_dropout=args.attention_dropout ,	activation_dropout=args.activation_dropout ,	activation_function="""relu""" ,	num_conv_layers=len(SCREAMING_SNAKE_CASE_						) ,	conv_channels=args.conv_channels ,	conv_kernel_sizes=SCREAMING_SNAKE_CASE_ ,	input_feat_per_channel=args.input_feat_per_channel ,	input_channels=args.input_channels ,	tie_word_embeddings=SCREAMING_SNAKE_CASE_ ,	num_beams=5 ,	max_length=2_00 ,	use_cache=SCREAMING_SNAKE_CASE_ ,	decoder_start_token_id=2 ,	early_stopping=SCREAMING_SNAKE_CASE_ ,	)

  lowercase_     =       SpeechaTextForConditionalGeneration(SCREAMING_SNAKE_CASE_						)
  lowercase_     ,	lowercase_     =       model.model.load_state_dict(SCREAMING_SNAKE_CASE_ ,	strict=SCREAMING_SNAKE_CASE_						)
  if len(SCREAMING_SNAKE_CASE_						) > 0 and not set(SCREAMING_SNAKE_CASE_						) <= {
      "encoder.embed_positions.weights",
      "decoder.embed_positions.weights",
  }:
    raise ValueError(
        """Only `encoder.embed_positions.weights` and `decoder.embed_positions.weights`  are allowed to be missing,"""
        f""" but all the following weights are missing {missing}"""						)

  if tie_embeds:
    lowercase_     =       make_linear_from_emb(model.model.decoder.embed_tokens						)
  else:
    lowercase_     =       lm_head_weights

  model.save_pretrained(SCREAMING_SNAKE_CASE_						)


if __name__ == "__main__":
       __snake_case						  =		argparse.ArgumentParser()
       # Required parameters
       parser.add_argument("""--fairseq_path""", type=str, help="""Path to the fairseq model (.pt) file.""")
       parser.add_argument("""--pytorch_dump_folder_path""", default=None, type=str, help="""Path to the output PyTorch model.""")
       __snake_case						  =		parser.parse_args()
       convert_fairseq_sat_checkpoint_to_tfms(args.fairseq_path, args.pytorch_dump_folder_path)

