'''simple docstring'''


import unittest

import numpy as np

from transformers import AlbertConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
       import jax.numpy as jnp

       from transformers.models.albert.modeling_flax_albert import (
           FlaxAlbertForMaskedLM,
           FlaxAlbertForMultipleChoice,
           FlaxAlbertForPreTraining,
           FlaxAlbertForQuestionAnswering,
           FlaxAlbertForSequenceClassification,
           FlaxAlbertForTokenClassification,
           FlaxAlbertModel,
       )





class       _a    (		unittest.TestCase      ):


  """simple docstring"""

  def __init__(		self      :       List[Any]   ,				lowercase_      :       Optional[int]   ,				lowercase_      :       str=13   ,				lowercase_      :       Any=7   ,				lowercase_      :       Any=True   ,				lowercase_      :       List[str]=True   ,				lowercase_      :       Tuple=True   ,				lowercase_      :       Optional[Any]=True   ,				lowercase_      :       Any=99   ,				lowercase_      :       int=32   ,				lowercase_      :       List[Any]=5   ,				lowercase_      :       Union[str, Any]=4   ,				lowercase_      :       Any=37   ,				lowercase_      :       Union[str, Any]="gelu"   ,				lowercase_      :       Any=0.1   ,				lowercase_      :       List[Any]=0.1   ,				lowercase_      :       Optional[int]=512   ,				lowercase_      :       Any=16   ,				lowercase_      :       Union[str, Any]=2   ,				lowercase_      :       int=0.0_2   ,				lowercase_      :       Union[str, Any]=4   ,				):







    '''simple docstring'''


    lowercase_     =       parent
    lowercase_     =       batch_size
    lowercase_     =       seq_length
    lowercase_     =       is_training
    lowercase_     =       use_attention_mask
    lowercase_     =       use_token_type_ids
    lowercase_     =       use_labels
    lowercase_     =       vocab_size
    lowercase_     =       hidden_size
    lowercase_     =       num_hidden_layers
    lowercase_     =       num_attention_heads
    lowercase_     =       intermediate_size
    lowercase_     =       hidden_act
    lowercase_     =       hidden_dropout_prob
    lowercase_     =       attention_probs_dropout_prob
    lowercase_     =       max_position_embeddings
    lowercase_     =       type_vocab_size
    lowercase_     =       type_sequence_label_size
    lowercase_     =       initializer_range
    lowercase_     =       num_choices




  def 		lowerCamelCase__  (		self      :       Dict      ):







    '''simple docstring'''


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

    lowercase_     =       None
    if self.use_attention_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_     =       AlbertConfig(
        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   ,				is_decoder=lowercase_   ,				initializer_range=self.initializer_range   ,				)

    return config, input_ids, token_type_ids, attention_mask





  def 		lowerCamelCase__  (		self      :       str      ):







    '''simple docstring'''


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






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


  """simple docstring"""

  A_    							=						(
      (
          FlaxAlbertModel,
          FlaxAlbertForPreTraining,
          FlaxAlbertForMaskedLM,
          FlaxAlbertForMultipleChoice,
          FlaxAlbertForQuestionAnswering,
          FlaxAlbertForSequenceClassification,
          FlaxAlbertForTokenClassification,
          FlaxAlbertForQuestionAnswering,
      )
      if is_flax_available()
      else ()
  )




  def 		lowerCamelCase__  (		self      :       Union[str, Any]      ):







    '''simple docstring'''


    lowercase_     =       FlaxAlbertModelTester(self      )





  @slow
  def 		lowerCamelCase__  (		self      :       Any      ):







    '''simple docstring'''


    for model_class_name in self.all_model_classes:
      lowercase_     =       model_class_name.from_pretrained("""albert-base-v2"""      )
      lowercase_     =       model(np.ones((1, 1)      )      )
      self.assertIsNotNone(lowercase_      )






@require_flax
class       _a    (		unittest.TestCase      ):


  """simple docstring"""

  @slow
  def 		lowerCamelCase__  (		self      :       int      ):







    '''simple docstring'''


    lowercase_     =       FlaxAlbertModel.from_pretrained("""albert-base-v2"""      )
    lowercase_     =       np.array([[0, 345, 232, 328, 740, 140, 1_695, 69, 6_078, 1_588, 2]]      )
    lowercase_     =       np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]      )
    lowercase_     =       model(lowercase_   ,				attention_mask=lowercase_      )[0]
    lowercase_     =       (1, 11, 768)
    self.assertEqual(output.shape   ,				lowercase_      )
    lowercase_     =       np.array(
        [[[-0.6_5_1_3, 1.5_0_3_5, -0.2_7_6_6], [-0.6_5_1_5, 1.5_0_4_6, -0.2_7_8_0], [-0.6_5_1_2, 1.5_0_4_9, -0.2_7_8_4]]]      )

    self.assertTrue(jnp.allclose(output[:, 1:4, 1:4]   ,				lowercase_   ,				atol=1e-4      )      )


'''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)

