

'''simple docstring'''





import argparse
import json
import os

import fairseq
import torch
from torch import nn

from transformers import (
    SpeechaTextaConfig,
    SpeechaTextaForCausalLM,
    SpeechaTextaTokenizer,
    SpeechEncoderDecoderConfig,
    SpeechEncoderDecoderModel,
    WavaVecaConfig,
    WavaVecaFeatureExtractor,
    WavaVecaModel,
    logging,
)


logging.set_verbosity_info()
A_	= logging.get_logger(__name__)

A_	= {
    "post_extract_proj": "feature_projection.projection",
    "encoder.pos_conv.0": "encoder.pos_conv_embed.conv",
    "self_attn.k_proj": "encoder.layers.*.attention.k_proj",
    "self_attn.v_proj": "encoder.layers.*.attention.v_proj",
    "self_attn.q_proj": "encoder.layers.*.attention.q_proj",
    "self_attn.out_proj": "encoder.layers.*.attention.out_proj",
    "self_attn_layer_norm": "encoder.layers.*.layer_norm",
    "fc1": "encoder.layers.*.feed_forward.intermediate_dense",
    "fc2": "encoder.layers.*.feed_forward.output_dense",
    "final_layer_norm": "encoder.layers.*.final_layer_norm",
    "encoder.layer_norm": "encoder.layer_norm",
    "w2v_model.layer_norm": "feature_projection.layer_norm",
    "quantizer.weight_proj": "quantizer.weight_proj",
    "quantizer.vars": "quantizer.codevectors",
    "project_q": "project_q",
    "final_proj": "project_hid",
    "w2v_encoder.proj": "lm_head",
    "mask_emb": "masked_spec_embed",
}
A_	= [
    "lm_head",
    "quantizer.weight_proj",
    "quantizer.codevectors",
    "project_q",
    "project_hid",
]







def   _UpperCamelCase	(		__UpperCamelCase    ,__UpperCamelCase    ,__UpperCamelCase    ,__UpperCamelCase    ,__UpperCamelCase	)       ->			Tuple:





					for attribute in key.split('.'	):
										lowerCamelCase_  				=      getattr(A__    ,A__	)

					if weight_type is not None:
										lowerCamelCase_  				=      getattr(A__    ,A__	).shape
					else:
										lowerCamelCase_  				=      hf_pointer.shape

					assert hf_shape == value.shape, (
					    f'''Shape of hf {key + "." + weight_type if weight_type is not None else ""} is {hf_shape}, but should be'''
					    f''' {value.shape} for {full_name}'''
					)

					if weight_type == "weight":
										lowerCamelCase_  				=      value
					elif weight_type == "weight_g":
										lowerCamelCase_  				=      value
					elif weight_type == "weight_v":
										lowerCamelCase_  				=      value
					elif weight_type == "bias":
										lowerCamelCase_  				=      value
					else:
										lowerCamelCase_  				=      value

					logger.info(f'''{key + "." + weight_type if weight_type is not None else ""} was initialized from {full_name}.'''	)







def   _UpperCamelCase	(		__UpperCamelCase    ,__UpperCamelCase	)       ->			Union[str, Any]:
					lowerCamelCase_  				=      []
					lowerCamelCase_  				=      fairseq_model.state_dict()

					lowerCamelCase_  				=      hf_model.feature_extractor

					# if encoder has different dim to decoder -> use proj_weight
					lowerCamelCase_  				=      None

					for name, value in fairseq_dict.items():
										lowerCamelCase_  				=      False
										if "conv_layers" in name:
															load_conv_layer(
															    A__    ,A__    ,A__    ,A__    ,hf_model.config.feat_extract_norm == 'group'    ,)
															lowerCamelCase_  				=      True
										elif name.split('.'	)[0] == "proj":
															lowerCamelCase_  				=      fairseq_model.proj
															lowerCamelCase_  				=      True
										else:
															for key, mapped_key in MAPPING.items():
																				if key in name or key.split('w2v_model.'	)[-1] == name.split('.'	)[0]:
																									lowerCamelCase_  				=      True
																									if "*" in mapped_key:
																														lowerCamelCase_  				=      name.split(A__	)[0].split('.'	)[-2]
																														lowerCamelCase_  				=      mapped_key.replace('*'    ,A__	)
																									if "weight_g" in name:
																														lowerCamelCase_  				=      """weight_g"""
																									elif "weight_v" in name:
																														lowerCamelCase_  				=      """weight_v"""
																									elif "bias" in name:
																														lowerCamelCase_  				=      """bias"""
																									elif "weight" in name:
																														lowerCamelCase_  				=      """weight"""
																									else:
																														lowerCamelCase_  				=      None
																									set_recursively(A__    ,A__    ,A__    ,A__    ,A__	)
																				continue
										if not is_used:
															unused_weights.append(A__	)

					logger.warning(f'''Unused weights: {unused_weights}'''	)

					return proj_weight







def   _UpperCamelCase	(		__UpperCamelCase    ,__UpperCamelCase    ,__UpperCamelCase    ,__UpperCamelCase    ,__UpperCamelCase	)       ->			List[str]:
					lowerCamelCase_  				=      full_name.split('conv_layers.'	)[-1]
					lowerCamelCase_  				=      name.split('.'	)
					lowerCamelCase_  				=      int(items[0]	)
					lowerCamelCase_  				=      int(items[1]	)

					if type_id == 0:
										if "bias" in name:
															assert value.shape == feature_extractor.conv_layers[layer_id].conv.bias.data.shape, (
															    f'''{full_name} has size {value.shape}, but'''
															    f''' {feature_extractor.conv_layers[layer_id].conv.bias.data.shape} was found.'''
															)
															lowerCamelCase_  				=      value
															logger.info(f'''Feat extract conv layer {layer_id} was initialized from {full_name}.'''	)
										elif "weight" in name:
															assert value.shape == feature_extractor.conv_layers[layer_id].conv.weight.data.shape, (
															    f'''{full_name} has size {value.shape}, but'''
															    f''' {feature_extractor.conv_layers[layer_id].conv.weight.data.shape} was found.'''
															)
															lowerCamelCase_  				=      value
															logger.info(f'''Feat extract conv layer {layer_id} was initialized from {full_name}.'''	)
					elif (type_id == 2 and not use_group_norm) or (type_id == 2 and layer_id == 0 and use_group_norm):
										if "bias" in name:
															assert value.shape == feature_extractor.conv_layers[layer_id].layer_norm.bias.data.shape, (
															    f'''{full_name} has size {value.shape}, but {feature_extractor[layer_id].layer_norm.bias.data.shape} was'''
															    " found."
															)
															lowerCamelCase_  				=      value
															logger.info(f'''Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.'''	)
										elif "weight" in name:
															assert value.shape == feature_extractor.conv_layers[layer_id].layer_norm.weight.data.shape, (
															    f'''{full_name} has size {value.shape}, but'''
															    f''' {feature_extractor[layer_id].layer_norm.weight.data.shape} was found.'''
															)
															lowerCamelCase_  				=      value
															logger.info(f'''Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.'''	)
					else:
										unused_weights.append(A__	)







def   _UpperCamelCase	(		__UpperCamelCase	)       ->			List[Any]:
					lowerCamelCase_  				=      emb.weight.shape
					lowerCamelCase_  				=      nn.Linear(A__    ,A__    ,bias=A__	)
					lowerCamelCase_  				=      emb.weight.data
					return lin_layer







def   _UpperCamelCase	(		__UpperCamelCase	)       ->			Optional[Any]:





					with open(A__    ,'r'    ,encoding='utf-8'	) as f:
										lowerCamelCase_  				=      f.readlines()
										lowerCamelCase_  				=      [line.split(' '	)[0] for line in lines]

					lowerCamelCase_  				=      len(A__	)

					lowerCamelCase_  				=      {
					    """<s>""": 0,
					    """<pad>""": 1,
					    """</s>""": 2,
					    """<unk>""": 3,
					}

					vocab_dict.update(dict(zip(A__    ,range(4    ,num_words + 4	)	)	)	)
					return vocab_dict







@torch.no_grad()
def   _UpperCamelCase	(		__UpperCamelCase    ,__UpperCamelCase    ,__UpperCamelCase    ,__UpperCamelCase    ,__UpperCamelCase    ,__UpperCamelCase    ,__UpperCamelCase    ,)       ->			Tuple:
					lowerCamelCase_  				=      WavaVecaConfig.from_pretrained(A__	)
					lowerCamelCase_  				=      SpeechaTextaConfig.from_pretrained(
					    A__    ,vocab_size=A__    ,decoder_layers=A__    ,do_stable_layer_norm=A__	)

					lowerCamelCase_  				=      WavaVecaFeatureExtractor(
					    feature_size=1    ,sampling_rate=1_60_00    ,padding_value=0    ,do_normalize=A__    ,return_attention_mask=A__    ,)

					lowerCamelCase_  				=      fairseq.checkpoint_utils.load_model_ensemble_and_task(
					    [checkpoint_path]    ,arg_overrides={'data': '/'.join(dict_path.split('/'	)[:-1]	)}	)
					lowerCamelCase_  				=      model[0].eval()

					# set weights for wav2vec2 encoder
					lowerCamelCase_  				=      WavaVecaModel(A__	)
					lowerCamelCase_  				=      recursively_load_weights_wavaveca(model.encoder    ,A__	)

					lowerCamelCase_  				=      SpeechaTextaForCausalLM(A__	)
					lowerCamelCase_  				=      hf_decoder.model.decoder.load_state_dict(model.decoder.state_dict()    ,strict=A__	)

					# set output linear layer
					unexpected_keys.remove('embed_out'	)
					lowerCamelCase_  				=      nn.Parameter(model.decoder.embed_out.detach()	)

					# layer norm is init to identity matrix so leaving it is fine
					logger.warning(f'''The following keys are missing when loading the decoder weights: {missing_keys}'''	)
					logger.warning(f'''The following keys are unexpected when loading the decoder weights: {unexpected_keys}'''	)

					lowerCamelCase_  				=      SpeechEncoderDecoderModel(encoder=A__    ,decoder=A__	)
					lowerCamelCase_  				=      False

					# add projection layer
					lowerCamelCase_  				=      nn.Parameter(projection_layer.weight	)
					lowerCamelCase_  				=      nn.Parameter(projection_layer.bias	)

					lowerCamelCase_  				=      create_vocab_dict(A__	)

					with open(os.path.join(A__    ,'vocab.json'	)    ,'w'	) as fp:
										json.dump(A__    ,A__	)

					lowerCamelCase_  				=      SpeechaTextaTokenizer(os.path.join(A__    ,'vocab.json'	)	)
					tokenizer.save_pretrained(A__	)

					lowerCamelCase_  				=      hf_wavavec.config.to_dict()
					lowerCamelCase_  				=      tokenizer.pad_token_id
					lowerCamelCase_  				=      tokenizer.bos_token_id
					lowerCamelCase_  				=      tokenizer.eos_token_id
					lowerCamelCase_  				=      """speech_to_text_2"""
					lowerCamelCase_  				=      """wav2vec2"""

					lowerCamelCase_  				=      SpeechEncoderDecoderConfig.from_dict(A__	)

					hf_wavavec.save_pretrained(A__	)
					feature_extractor.save_pretrained(A__	)


if __name__ == "__main__":
				A_	= argparse.ArgumentParser()
				parser.add_argument("--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model.")
				parser.add_argument("--checkpoint_path", default=None, type=str, help="Path to fairseq checkpoint")
				parser.add_argument("--dict_path", default=None, type=str, help="Path to dict of fine-tuned model")
				parser.add_argument(
				    "--encoder_config_path",
				    default="facebook/wav2vec2-large-lv60",
				    type=str,
				    help="Path to hf encoder wav2vec2 checkpoint config",
				)
				parser.add_argument(
				    "--decoder_config_path",
				    default="facebook/s2t-small-mustc-en-fr-st",
				    type=str,
				    help="Path to hf decoder s2t checkpoint config",
				)
				parser.add_argument("--vocab_size", default=10_224, type=int, help="Vocab size of decoder")
				parser.add_argument("--num_decoder_layers", default=7, type=int, help="Number of decoder layers")

				A_	= parser.parse_args()
				convert_wavaveca_checkpoint(
				    args.checkpoint_path,
				    args.pytorch_dump_folder_path,
				    args.dict_path,
				    encoder_config_path=args.encoder_config_path,
				    decoder_config_path=args.decoder_config_path,
				    vocab_size=args.vocab_size,
				    num_decoder_layers=args.num_decoder_layers,
				)







'''simple docstring'''






import inspect
import tempfile
import unittest

from huggingface_hub import hf_hub_download

from transformers import is_torch_available
from transformers.testing_utils import is_flaky, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


__A   :    Dict											= 1E-4

if is_torch_available():
  import torch

  from transformers import AutoformerConfig, AutoformerForPrediction, AutoformerModel
  from transformers.models.autoformer.modeling_autoformer import AutoformerDecoder, AutoformerEncoder






@require_torch
class   __snake_case			:







    """simple docstring"""



    def __init__(  self       :    Tuple       ,		lowerCamelCase       :    Optional[Any]       ,		lowerCamelCase       :    List[str]=16       ,		lowerCamelCase       :    Dict=13       ,		lowerCamelCase       :    List[Any]=7       ,		lowerCamelCase       :    Optional[Any]=14       ,		lowerCamelCase       :    int=10       ,		lowerCamelCase       :    str=19       ,		lowerCamelCase       :    int=5       ,		lowerCamelCase       :    int=4       ,		lowerCamelCase       :    Any=True       ,		lowerCamelCase       :    Optional[Any]=16       ,		lowerCamelCase       :    Any=2       ,		lowerCamelCase       :    str=4       ,		lowerCamelCase       :    str=4       ,		lowerCamelCase       :    Tuple="gelu"       ,		lowerCamelCase       :    Optional[int]=0.1       ,		lowerCamelCase       :    int=0.1       ,		lowerCamelCase       :    Union[str, Any]=[1, 2, 3, 4, 5]       ,		lowerCamelCase       :    Dict=25       ,		lowerCamelCase       :    str=5       ,		)  ->    Optional[Any]:
      lowerCAmelCase_     : Optional[int]     =  d_model
      lowerCAmelCase_     : Tuple     =  parent
      lowerCAmelCase_     : int     =  batch_size
      lowerCAmelCase_     : List[Any]     =  prediction_length
      lowerCAmelCase_     : Tuple     =  context_length
      lowerCAmelCase_     : Tuple     =  cardinality
      lowerCAmelCase_     : Union[str, Any]     =  num_time_features
      lowerCAmelCase_     : List[str]     =  lags_sequence
      lowerCAmelCase_     : Dict     =  embedding_dimension
      lowerCAmelCase_     : Tuple     =  is_training
      lowerCAmelCase_     : Dict     =  hidden_size
      lowerCAmelCase_     : Optional[Any]     =  num_hidden_layers
      lowerCAmelCase_     : Any     =  num_attention_heads
      lowerCAmelCase_     : Dict     =  intermediate_size
      lowerCAmelCase_     : List[str]     =  hidden_act
      lowerCAmelCase_     : int     =  hidden_dropout_prob
      lowerCAmelCase_     : Any     =  attention_probs_dropout_prob

      lowerCAmelCase_     : str     =  context_length
      lowerCAmelCase_     : int     =  prediction_length + label_length
      lowerCAmelCase_     : int     =  label_length

      lowerCAmelCase_     : str     =  moving_average
      lowerCAmelCase_     : Optional[int]     =  autocorrelation_factor

    def __lowercase  (  self       :    int						)  ->    List[str]:
      return AutoformerConfig(
          d_model=self.d_model       ,		encoder_layers=self.num_hidden_layers       ,		decoder_layers=self.num_hidden_layers       ,		encoder_attention_heads=self.num_attention_heads       ,		decoder_attention_heads=self.num_attention_heads       ,		encoder_ffn_dim=self.intermediate_size       ,		decoder_ffn_dim=self.intermediate_size       ,		dropout=self.hidden_dropout_prob       ,		attention_dropout=self.attention_probs_dropout_prob       ,		prediction_length=self.prediction_length       ,		context_length=self.context_length       ,		label_length=self.label_length       ,		lags_sequence=self.lags_sequence       ,		num_time_features=self.num_time_features       ,		num_static_categorical_features=1       ,		cardinality=[self.cardinality]       ,		embedding_dimension=[self.embedding_dimension]       ,		moving_average=self.moving_average       ,		)

    def __lowercase  (  self       :    Any       ,		lowerCamelCase       :    Dict						)  ->    List[str]:
      lowerCAmelCase_     : Union[str, Any]     =  config.context_length + max(config.lags_sequence						)

      lowerCAmelCase_     : Any     =  ids_tensor([self.batch_size, 1]       ,		config.cardinality[0]						)
      lowerCAmelCase_     : Any     =  floats_tensor([self.batch_size, _past_length, config.num_time_features]						)
      lowerCAmelCase_     : Tuple     =  floats_tensor([self.batch_size, _past_length]						)
      lowerCAmelCase_     : List[str]     =  floats_tensor([self.batch_size, _past_length]						) > 0.5

      # decoder inputs
      lowerCAmelCase_     : List[str]     =  floats_tensor([self.batch_size, config.prediction_length, config.num_time_features]						)
      lowerCAmelCase_     : Optional[int]     =  floats_tensor([self.batch_size, config.prediction_length]						)

      lowerCAmelCase_     : List[str]     =  {
          """past_values""": past_values,
          """static_categorical_features""": static_categorical_features,
          """past_time_features""": past_time_features,
          """past_observed_mask""": past_observed_mask,
          """future_time_features""": future_time_features,
          """future_values""": future_values,
      }
      return inputs_dict

    def __lowercase  (  self       :    Tuple						)  ->    Optional[Any]:
      lowerCAmelCase_     : List[Any]     =  self.get_config()
      lowerCAmelCase_     : str     =  self.prepare_autoformer_inputs_dict(lowerCamelCase						)
      return config, inputs_dict

    def __lowercase  (  self       :    Optional[int]						)  ->    Optional[Any]:
      lowerCAmelCase_,				lowerCAmelCase_     : List[str]     =  self.prepare_config_and_inputs()
      return config, inputs_dict





    def __lowercase  (  self       :    str       ,		lowerCamelCase       :    Any       ,		lowerCamelCase       :    Dict						)  ->    Union[str, Any]:
      lowerCAmelCase_     : int     =  AutoformerModel(config=lowerCamelCase						).to(lowerCamelCase						).eval()
      lowerCAmelCase_     : List[str]     =  model(**lowerCamelCase						)

      lowerCAmelCase_     : Any     =  outputs.encoder_last_hidden_state
      lowerCAmelCase_     : List[str]     =  outputs.last_hidden_state

      with tempfile.TemporaryDirectory() as tmpdirname:
        lowerCAmelCase_     : Optional[int]     =  model.get_encoder()
        encoder.save_pretrained(lowerCamelCase						)
        lowerCAmelCase_     : List[Any]     =  AutoformerEncoder.from_pretrained(lowerCamelCase						).to(lowerCamelCase						)

      lowerCAmelCase_,				lowerCAmelCase_,				lowerCAmelCase_,				lowerCAmelCase_,				lowerCAmelCase_     : Dict     =  model.create_network_inputs(**lowerCamelCase						)
      lowerCAmelCase_,				lowerCAmelCase_     : Optional[Any]     =  model.decomposition_layer(transformer_inputs[:, : config.context_length, ...]						)

      lowerCAmelCase_     : Tuple     =  torch.cat(
          (transformer_inputs[:, : config.context_length, ...], feature[:, : config.context_length, ...])       ,		dim=-1       ,		)
      lowerCAmelCase_     : Union[str, Any]     =  encoder(inputs_embeds=lowerCamelCase						)[0]
      self.parent.assertTrue((encoder_last_hidden_state_a - encoder_last_hidden_state).abs().max().item() < 1E-3						)

      lowerCAmelCase_     : List[Any]     =  (
          torch.mean(transformer_inputs[:, : config.context_length, ...]       ,		dim=1						)
          .unsqueeze(1						)
          .repeat(1       ,		config.prediction_length       ,		1						)
      )
      lowerCAmelCase_     : Tuple     =  torch.zeros(
          [transformer_inputs.shape[0], config.prediction_length, transformer_inputs.shape[2]]       ,		device=enc_input.device       ,		)

      lowerCAmelCase_     : Union[str, Any]     =  torch.cat(
          (
              torch.cat((seasonal_input[:, -config.label_length :, ...], zeros)       ,		dim=1						),
              feature[:, config.context_length - config.label_length :, ...],
          )       ,		dim=-1       ,		)
      lowerCAmelCase_     : int     =  torch.cat(
          (
              torch.cat((trend_input[:, -config.label_length :, ...], mean)       ,		dim=1						),
              feature[:, config.context_length - config.label_length :, ...],
          )       ,		dim=-1       ,		)

      with tempfile.TemporaryDirectory() as tmpdirname:
        lowerCAmelCase_     : Any     =  model.get_decoder()
        decoder.save_pretrained(lowerCamelCase						)
        lowerCAmelCase_     : Optional[int]     =  AutoformerDecoder.from_pretrained(lowerCamelCase						).to(lowerCamelCase						)

      lowerCAmelCase_     : int     =  decoder(
          trend=lowerCamelCase       ,		inputs_embeds=lowerCamelCase       ,		encoder_hidden_states=lowerCamelCase       ,		)[0]

      self.parent.assertTrue((last_hidden_state_a - last_hidden_state).abs().max().item() < 1E-3						)







@require_torch
class   __snake_case			(  _SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,unittest.TestCase):







    """simple docstring"""



    lowercase  					=   (AutoformerModel, AutoformerForPrediction) if is_torch_available() else ()
    lowercase  					=   (AutoformerForPrediction,) if is_torch_available() else ()
    lowercase  					=   {'feature-extraction': AutoformerModel} if is_torch_available() else {}
    lowercase  					=   False
    lowercase  					=   False
    lowercase  					=   False
    lowercase  					=   False
    lowercase  					=   False
    lowercase  					=   False

    def __lowercase  (  self       :    Optional[Any]						)  ->    List[str]:
      lowerCAmelCase_     : str     =  AutoformerModelTester(self						)
      lowerCAmelCase_     : Union[str, Any]     =  ConfigTester(self       ,		config_class=lowerCamelCase       ,		has_text_modality=lowerCamelCase						)

    def __lowercase  (  self       :    Dict						)  ->    List[str]:
      self.config_tester.run_common_tests()

    def __lowercase  (  self       :    Optional[Any]						)  ->    Optional[int]:
      lowerCAmelCase_,				lowerCAmelCase_     : Tuple     =  self.model_tester.prepare_config_and_inputs()
      for model_class in self.all_model_classes:
        lowerCAmelCase_     : int     =  model_class(lowerCamelCase						)

        with tempfile.TemporaryDirectory() as tmpdirname:
          model.save_pretrained(lowerCamelCase						)
          lowerCAmelCase_,				lowerCAmelCase_     : Union[str, Any]     =  model_class.from_pretrained(lowerCamelCase       ,		output_loading_info=lowerCamelCase						)
        self.assertEqual(info["""missing_keys"""]       ,		[]						)

    def __lowercase  (  self       :    str						)  ->    str:
      lowerCAmelCase_     : Union[str, Any]     =  self.model_tester.prepare_config_and_inputs_for_common()
      self.model_tester.check_encoder_decoder_model_standalone(*lowerCamelCase						)

    @unittest.skip(reason="""Model has no tokens embeddings"""						)
    def __lowercase  (  self       :    List[Any]						)  ->    Dict:
      pass

    def __lowercase  (  self       :    Optional[Any]						)  ->    Tuple:
      lowerCAmelCase_     : Dict     =  inspect.signature(getattr(lowerCamelCase       ,		"""forward"""						)						)
      # The main input is the name of the argument after `self`
      lowerCAmelCase_     : Optional[int]     =  list(model_signature.parameters.keys()						)[1]
      self.assertEqual(AutoformerModel.main_input_name       ,		lowerCamelCase						)

    def __lowercase  (  self       :    str						)  ->    List[str]:
      lowerCAmelCase_,				lowerCAmelCase_     : Dict     =  self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
        lowerCAmelCase_     : Dict     =  model_class(lowerCamelCase						)
        lowerCAmelCase_     : Optional[int]     =  inspect.signature(model.forward						)
        # signature.parameters is an OrderedDict => so arg_names order is deterministic
        lowerCAmelCase_     : str     =  [*signature.parameters.keys()]

        lowerCAmelCase_     : Any     =  [
            """past_values""",
            """past_time_features""",
            """past_observed_mask""",
            """static_categorical_features""",
            """static_real_features""",
            """future_values""",
            """future_time_features""",
        ]

        if model.__class__.__name__ in ["AutoformerForPrediction"]:
          expected_arg_names.append("""future_observed_mask"""						)

        expected_arg_names.extend(
            [
                """decoder_attention_mask""",
                """head_mask""",
                """decoder_head_mask""",
                """cross_attn_head_mask""",
                """encoder_outputs""",
                """past_key_values""",
                """output_hidden_states""",
                """output_attentions""",
                """use_cache""",
                """return_dict""",
            ]						)

        self.assertListEqual(arg_names[: len(lowerCamelCase						)]       ,		lowerCamelCase						)

    def __lowercase  (  self       :    Optional[int]						)  ->    Optional[int]:
      lowerCAmelCase_,				lowerCAmelCase_     : int     =  self.model_tester.prepare_config_and_inputs_for_common()
      lowerCAmelCase_     : Tuple     =  True

      lowerCAmelCase_     : Tuple     =  getattr(self.model_tester       ,		"""seq_length"""       ,		lowerCamelCase						)
      lowerCAmelCase_     : Union[str, Any]     =  getattr(self.model_tester       ,		"""decoder_seq_length"""       ,		lowerCamelCase						)
      lowerCAmelCase_     : Tuple     =  getattr(self.model_tester       ,		"""encoder_seq_length"""       ,		lowerCamelCase						)
      lowerCAmelCase_     : Optional[int]     =  getattr(self.model_tester       ,		"""d_model"""       ,		lowerCamelCase						)
      lowerCAmelCase_     : Any     =  getattr(self.model_tester       ,		"""num_attention_heads"""       ,		lowerCamelCase						)
      lowerCAmelCase_     : Optional[Any]     =  d_model // num_attention_heads

      for model_class in self.all_model_classes:
        lowerCAmelCase_     : List[Any]     =  True
        lowerCAmelCase_     : Optional[int]     =  False
        lowerCAmelCase_     : Optional[int]     =  True
        lowerCAmelCase_     : List[Any]     =  model_class(lowerCamelCase						)
        model.to(lowerCamelCase						)
        model.eval()
        with torch.no_grad():
          lowerCAmelCase_     : str     =  model(**self._prepare_for_class(lowerCamelCase       ,		lowerCamelCase						)						)
        lowerCAmelCase_     : Any     =  outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions
        self.assertEqual(len(lowerCamelCase						)       ,		self.model_tester.num_hidden_layers						)

        # check that output_attentions also work using config
        del inputs_dict["output_attentions"]
        lowerCAmelCase_     : str     =  True
        lowerCAmelCase_     : Any     =  model_class(lowerCamelCase						)
        model.to(lowerCamelCase						)
        model.eval()
        with torch.no_grad():
          lowerCAmelCase_     : List[Any]     =  model(**self._prepare_for_class(lowerCamelCase       ,		lowerCamelCase						)						)
        lowerCAmelCase_     : Optional[Any]     =  outputs.encoder_attentions
        self.assertEqual(len(lowerCamelCase						)       ,		self.model_tester.num_hidden_layers						)

        self.assertListEqual(
            list(attentions[0].shape[-3:]						)       ,		[self.model_tester.num_attention_heads, encoder_seq_length, dim]       ,		)
        lowerCAmelCase_     : Any     =  len(lowerCamelCase						)

        lowerCAmelCase_     : Any     =  7

        if "last_hidden_state" in outputs:
          correct_outlen += 1

        if "trend" in outputs:
          correct_outlen += 1

        if "past_key_values" in outputs:
          correct_outlen += 1  # past_key_values have been returned

        if "loss" in outputs:
          correct_outlen += 1

        if "params" in outputs:
          correct_outlen += 1

        self.assertEqual(lowerCamelCase       ,		lowerCamelCase						)

        # decoder attentions
        lowerCAmelCase_     : Any     =  outputs.decoder_attentions
        self.assertIsInstance(lowerCamelCase       ,		(list, tuple)						)
        self.assertEqual(len(lowerCamelCase						)       ,		self.model_tester.num_hidden_layers						)
        self.assertListEqual(
            list(decoder_attentions[0].shape[-3:]						)       ,		[self.model_tester.num_attention_heads, decoder_seq_length, dim]       ,		)

        # cross attentions
        lowerCAmelCase_     : int     =  outputs.cross_attentions
        self.assertIsInstance(lowerCamelCase       ,		(list, tuple)						)
        self.assertEqual(len(lowerCamelCase						)       ,		self.model_tester.num_hidden_layers						)
        self.assertListEqual(
            list(cross_attentions[0].shape[-3:]						)       ,		[self.model_tester.num_attention_heads, decoder_seq_length, dim]       ,		)

      # Check attention is always last and order is fine
      lowerCAmelCase_     : Union[str, Any]     =  True
      lowerCAmelCase_     : List[str]     =  True
      lowerCAmelCase_     : List[str]     =  model_class(lowerCamelCase						)
      model.to(lowerCamelCase						)
      model.eval()
      with torch.no_grad():
        lowerCAmelCase_     : str     =  model(**self._prepare_for_class(lowerCamelCase       ,		lowerCamelCase						)						)

      self.assertEqual(out_len + 2       ,		len(lowerCamelCase						)						)

      lowerCAmelCase_     : int     =  outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions

      self.assertEqual(len(lowerCamelCase						)       ,		self.model_tester.num_hidden_layers						)
      self.assertListEqual(
          list(self_attentions[0].shape[-3:]						)       ,		[self.model_tester.num_attention_heads, encoder_seq_length, dim]       ,		)





    @is_flaky()
    def __lowercase  (  self       :    Any						)  ->    Dict:
      super().test_retain_grad_hidden_states_attentions()







def 							UpperCamelCase_     (       A__   :  Optional[Any]="train-batch.pt"   ):
  '''simple docstring'''





  lowerCAmelCase_     : Dict     =  hf_hub_download(repo_id="""hf-internal-testing/tourism-monthly-batch"""				,       filename=A__				,       repo_type="""dataset"""   )
  lowerCAmelCase_     : int     =  torch.load(A__				,       map_location=A__   )
  return batch






@require_torch
@slow
class   __snake_case			(  unittest.TestCase):







    """simple docstring"""



    def __lowercase  (  self       :    int						)  ->    Optional[Any]:
      lowerCAmelCase_     : Optional[int]     =  AutoformerModel.from_pretrained("""huggingface/autoformer-tourism-monthly"""						).to(lowerCamelCase						)
      lowerCAmelCase_     : Optional[Any]     =  prepare_batch()

      with torch.no_grad():
        lowerCAmelCase_     : Optional[int]     =  model(
            past_values=batch["""past_values"""]       ,		past_time_features=batch["""past_time_features"""]       ,		past_observed_mask=batch["""past_observed_mask"""]       ,		static_categorical_features=batch["""static_categorical_features"""]       ,		future_values=batch["""future_values"""]       ,		future_time_features=batch["""future_time_features"""]       ,		)[0]

      lowerCAmelCase_     : List[Any]     =  torch.Size(
          (64, model.config.prediction_length + model.config.label_length, model.config.feature_size)						)
      self.assertEqual(output.shape       ,		lowerCamelCase						)

      lowerCAmelCase_     : Dict     =  torch.tensor(
          [[0.3_593, -1.3_398, 0.6_330], [0.2_279, 1.5_396, -0.1_792], [0.0_450, 1.3_225, -0.2_335]]       ,		device=lowerCamelCase						)
      self.assertTrue(torch.allclose(output[0, :3, :3]       ,		lowerCamelCase       ,		atol=lowerCamelCase						)						)

    def __lowercase  (  self       :    Union[str, Any]						)  ->    Any:
      lowerCAmelCase_     : List[str]     =  AutoformerForPrediction.from_pretrained("""huggingface/autoformer-tourism-monthly"""						).to(lowerCamelCase						)
      lowerCAmelCase_     : str     =  prepare_batch("""val-batch.pt"""						)
      with torch.no_grad():
        lowerCAmelCase_     : List[str]     =  model(
            past_values=batch["""past_values"""]       ,		past_time_features=batch["""past_time_features"""]       ,		past_observed_mask=batch["""past_observed_mask"""]       ,		static_categorical_features=batch["""static_categorical_features"""]       ,		).encoder_last_hidden_state
      lowerCAmelCase_     : List[Any]     =  torch.Size((64, model.config.context_length, model.config.d_model)						)
      self.assertEqual(output.shape       ,		lowerCamelCase						)

      lowerCAmelCase_     : List[Any]     =  torch.tensor(
          [[-0.0_734, -0.9_036, 0.8_358], [4.7_186, 2.4_113, 1.9_581], [1.7_953, 2.3_558, 1.2_970]]       ,		device=lowerCamelCase						)
      self.assertTrue(torch.allclose(output[0, :3, :3]       ,		lowerCamelCase       ,		atol=lowerCamelCase						)						)





    def __lowercase  (  self       :    Any						)  ->    Optional[Any]:
      lowerCAmelCase_     : List[str]     =  AutoformerForPrediction.from_pretrained("""huggingface/autoformer-tourism-monthly"""						).to(lowerCamelCase						)
      lowerCAmelCase_     : int     =  prepare_batch("""val-batch.pt"""						)
      with torch.no_grad():
        lowerCAmelCase_     : Union[str, Any]     =  model.generate(
            static_categorical_features=batch["""static_categorical_features"""]       ,		past_time_features=batch["""past_time_features"""]       ,		past_values=batch["""past_values"""]       ,		future_time_features=batch["""future_time_features"""]       ,		past_observed_mask=batch["""past_observed_mask"""]       ,		)
      lowerCAmelCase_     : Any     =  torch.Size((64, model.config.num_parallel_samples, model.config.prediction_length)						)
      self.assertEqual(outputs.sequences.shape       ,		lowerCamelCase						)

      lowerCAmelCase_     : List[Any]     =  torch.tensor([3_130.6_763, 4_056.5_293, 7_053.0_786]       ,		device=lowerCamelCase						)
      lowerCAmelCase_     : str     =  outputs.sequences.mean(dim=1						)
      self.assertTrue(torch.allclose(mean_prediction[0, -3:]       ,		lowerCamelCase       ,		rtol=1E-1						)						)



