





"""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


SCREAMING_SNAKE_CASE__					=   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 		lowercase					:





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

            lowerCAmelCase  						=       context_length
            lowerCAmelCase  						=       prediction_length + label_length
            lowerCAmelCase  						=       label_length

            lowerCAmelCase  						=       moving_average
            lowerCAmelCase  						=       autocorrelation_factor





      def     _snake_case			(       self )		->     Optional[int]:
            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     _snake_case			(       self      ,						lowercase )		->     Optional[int]:
            lowerCAmelCase  						=       config.context_length + max(config.lags_sequence )

            lowerCAmelCase  						=       ids_tensor([self.batch_size, 1]      ,						config.cardinality[0] )
            lowerCAmelCase  						=       floats_tensor([self.batch_size, _past_length, config.num_time_features] )
            lowerCAmelCase  						=       floats_tensor([self.batch_size, _past_length] )
            lowerCAmelCase  						=       floats_tensor([self.batch_size, _past_length] ) > 0.5

            # decoder inputs
            lowerCAmelCase  						=       floats_tensor([self.batch_size, config.prediction_length, config.num_time_features] )
            lowerCAmelCase  						=       floats_tensor([self.batch_size, config.prediction_length] )

            lowerCAmelCase  						=       {
                """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     _snake_case			(       self )		->     str:
            lowerCAmelCase  						=       self.get_config()
            lowerCAmelCase  						=       self.prepare_autoformer_inputs_dict(lowercase )
            return config, inputs_dict





      def     _snake_case			(       self )		->     Any:
            lowerCAmelCase						,		lowerCAmelCase  						=       self.prepare_config_and_inputs()
            return config, inputs_dict





      def     _snake_case			(       self      ,						lowercase      ,						lowercase )		->     Union[str, Any]:
            lowerCAmelCase  						=       AutoformerModel(config=lowercase ).to(lowercase ).eval()
            lowerCAmelCase  						=       model(**lowercase )

            lowerCAmelCase  						=       outputs.encoder_last_hidden_state
            lowerCAmelCase  						=       outputs.last_hidden_state

            with tempfile.TemporaryDirectory() as tmpdirname:
                  lowerCAmelCase  						=       model.get_encoder()
                  encoder.save_pretrained(lowercase )
                  lowerCAmelCase  						=       AutoformerEncoder.from_pretrained(lowercase ).to(lowercase )

            lowerCAmelCase						,		lowerCAmelCase						,		lowerCAmelCase						,		lowerCAmelCase						,		lowerCAmelCase  						=       model.create_network_inputs(**lowercase )
            lowerCAmelCase						,		lowerCAmelCase  						=       model.decomposition_layer(transformer_inputs[:, : config.context_length, ...] )

            lowerCAmelCase  						=       torch.cat(
                (transformer_inputs[:, : config.context_length, ...], feature[:, : config.context_length, ...])      ,						dim=-1      ,						)
            lowerCAmelCase  						=       encoder(inputs_embeds=lowercase )[0]
            self.parent.assertTrue((encoder_last_hidden_state_a - encoder_last_hidden_state).abs().max().item() < 1e-3 )

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

            lowerCAmelCase  						=       torch.cat(
                (
                    torch.cat((seasonal_input[:, -config.label_length :, ...], zeros)      ,						dim=1 ),
                    feature[:, config.context_length - config.label_length :, ...],
                )      ,						dim=-1      ,						)
            lowerCAmelCase  						=       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  						=       model.get_decoder()
                  decoder.save_pretrained(lowercase )
                  lowerCAmelCase  						=       AutoformerDecoder.from_pretrained(lowercase ).to(lowercase )

            lowerCAmelCase  						=       decoder(
                trend=lowercase      ,						inputs_embeds=lowercase      ,						encoder_hidden_states=lowercase      ,						)[0]

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







@require_torch
class 		lowercase					(					_UpperCAmelCase     ,  _UpperCAmelCase     ,  unittest.TestCase					):
      _SCREAMING_SNAKE_CASE					=      (AutoformerModel, AutoformerForPrediction) if is_torch_available() else ()
      _SCREAMING_SNAKE_CASE					=      (AutoformerForPrediction,) if is_torch_available() else ()
      _SCREAMING_SNAKE_CASE					=      {'feature-extraction': AutoformerModel} if is_torch_available() else {}
      _SCREAMING_SNAKE_CASE					=      False
      _SCREAMING_SNAKE_CASE					=      False
      _SCREAMING_SNAKE_CASE					=      False
      _SCREAMING_SNAKE_CASE					=      False
      _SCREAMING_SNAKE_CASE					=      False
      _SCREAMING_SNAKE_CASE					=      False





      def     _snake_case			(       self )		->     Union[str, Any]:
            lowerCAmelCase  						=       AutoformerModelTester(self )
            lowerCAmelCase  						=       ConfigTester(self      ,						config_class=lowercase      ,						has_text_modality=lowercase )





      def     _snake_case			(       self )		->     Optional[Any]:
            self.config_tester.run_common_tests()





      def     _snake_case			(       self )		->     str:
            lowerCAmelCase						,		lowerCAmelCase  						=       self.model_tester.prepare_config_and_inputs()
            for model_class in self.all_model_classes:
                  lowerCAmelCase  						=       model_class(lowercase )

                  with tempfile.TemporaryDirectory() as tmpdirname:
                        model.save_pretrained(lowercase )
                        lowerCAmelCase						,		lowerCAmelCase  						=       model_class.from_pretrained(lowercase      ,						output_loading_info=lowercase )
                  self.assertEqual(info["""missing_keys"""]      ,						[] )





      def     _snake_case			(       self )		->     Any:
            lowerCAmelCase  						=       self.model_tester.prepare_config_and_inputs_for_common()
            self.model_tester.check_encoder_decoder_model_standalone(*lowercase )





      @unittest.skip(reason="""Model has no tokens embeddings""" )
      def     _snake_case			(       self )		->     str:
            pass





      def     _snake_case			(       self )		->     int:
            lowerCAmelCase  						=       inspect.signature(getattr(lowercase      ,						"""forward""" ) )
            # The main input is the name of the argument after `self`
            lowerCAmelCase  						=       list(model_signature.parameters.keys() )[1]
            self.assertEqual(AutoformerModel.main_input_name      ,						lowercase )





      def     _snake_case			(       self )		->     int:
            lowerCAmelCase						,		lowerCAmelCase  						=       self.model_tester.prepare_config_and_inputs_for_common()

            for model_class in self.all_model_classes:
                  lowerCAmelCase  						=       model_class(lowercase )
                  lowerCAmelCase  						=       inspect.signature(model.forward )
                  # signature.parameters is an OrderedDict => so arg_names order is deterministic
                  lowerCAmelCase  						=       [*signature.parameters.keys()]

                  lowerCAmelCase  						=       [
                      """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(lowercase )]      ,						lowercase )





      def     _snake_case			(       self )		->     Optional[Any]:
            lowerCAmelCase						,		lowerCAmelCase  						=       self.model_tester.prepare_config_and_inputs_for_common()
            lowerCAmelCase  						=       True

            lowerCAmelCase  						=       getattr(self.model_tester      ,						"""seq_length"""      ,						lowercase )
            lowerCAmelCase  						=       getattr(self.model_tester      ,						"""decoder_seq_length"""      ,						lowercase )
            lowerCAmelCase  						=       getattr(self.model_tester      ,						"""encoder_seq_length"""      ,						lowercase )
            lowerCAmelCase  						=       getattr(self.model_tester      ,						"""d_model"""      ,						lowercase )
            lowerCAmelCase  						=       getattr(self.model_tester      ,						"""num_attention_heads"""      ,						lowercase )
            lowerCAmelCase  						=       d_model // num_attention_heads

            for model_class in self.all_model_classes:
                  lowerCAmelCase  						=       True
                  lowerCAmelCase  						=       False
                  lowerCAmelCase  						=       True
                  lowerCAmelCase  						=       model_class(lowercase )
                  model.to(lowercase )
                  model.eval()
                  with torch.no_grad():
                        lowerCAmelCase  						=       model(**self._prepare_for_class(lowercase      ,						lowercase ) )
                  lowerCAmelCase  						=       outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions
                  self.assertEqual(len(lowercase )      ,						self.model_tester.num_hidden_layers )

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

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

                  lowerCAmelCase  						=       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(lowercase      ,						lowercase )

                  # decoder attentions
                  lowerCAmelCase  						=       outputs.decoder_attentions
                  self.assertIsInstance(lowercase      ,						(list, tuple) )
                  self.assertEqual(len(lowercase )      ,						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  						=       outputs.cross_attentions
                  self.assertIsInstance(lowercase      ,						(list, tuple) )
                  self.assertEqual(len(lowercase )      ,						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  						=       True
            lowerCAmelCase  						=       True
            lowerCAmelCase  						=       model_class(lowercase )
            model.to(lowercase )
            model.eval()
            with torch.no_grad():
                  lowerCAmelCase  						=       model(**self._prepare_for_class(lowercase      ,						lowercase ) )

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

            lowerCAmelCase  						=       outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions

            self.assertEqual(len(lowercase )      ,						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     _snake_case			(       self )		->     Any:
            super().test_retain_grad_hidden_states_attentions()







def 			UpperCAmelCase__							(  SCREAMING_SNAKE_CASE     :						List[str]="train-batch.pt"			):







      '''simple docstring'''




      lowerCAmelCase  						=       hf_hub_download(repo_id="""hf-internal-testing/tourism-monthly-batch"""							,						filename=SCREAMING_SNAKE_CASE							,						repo_type="""dataset"""			)
      lowerCAmelCase  						=       torch.load(SCREAMING_SNAKE_CASE							,						map_location=SCREAMING_SNAKE_CASE			)
      return batch






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





      def     _snake_case			(       self )		->     Optional[Any]:
            lowerCAmelCase  						=       AutoformerModel.from_pretrained("""huggingface/autoformer-tourism-monthly""" ).to(lowercase )
            lowerCAmelCase  						=       prepare_batch()

            with torch.no_grad():
                  lowerCAmelCase  						=       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  						=       torch.Size(
                (64, model.config.prediction_length + model.config.label_length, model.config.feature_size) )
            self.assertEqual(output.shape      ,						lowercase )

            lowerCAmelCase  						=       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=lowercase )
            self.assertTrue(torch.allclose(output[0, :3, :3]      ,						lowercase      ,						atol=lowercase ) )





      def     _snake_case			(       self )		->     List[Any]:
            lowerCAmelCase  						=       AutoformerForPrediction.from_pretrained("""huggingface/autoformer-tourism-monthly""" ).to(lowercase )
            lowerCAmelCase  						=       prepare_batch("""val-batch.pt""" )
            with torch.no_grad():
                  lowerCAmelCase  						=       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  						=       torch.Size((64, model.config.context_length, model.config.d_model) )
            self.assertEqual(output.shape      ,						lowercase )

            lowerCAmelCase  						=       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=lowercase )
            self.assertTrue(torch.allclose(output[0, :3, :3]      ,						lowercase      ,						atol=lowercase ) )





      def     _snake_case			(       self )		->     Union[str, Any]:
            lowerCAmelCase  						=       AutoformerForPrediction.from_pretrained("""huggingface/autoformer-tourism-monthly""" ).to(lowercase )
            lowerCAmelCase  						=       prepare_batch("""val-batch.pt""" )
            with torch.no_grad():
                  lowerCAmelCase  						=       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  						=       torch.Size((64, model.config.num_parallel_samples, model.config.prediction_length) )
            self.assertEqual(outputs.sequences.shape      ,						lowercase )

            lowerCAmelCase  						=       torch.tensor([3_130.6_763, 4_056.5_293, 7_053.0_786]      ,						device=lowercase )
            lowerCAmelCase  						=       outputs.sequences.mean(dim=1 )
            self.assertTrue(torch.allclose(mean_prediction[0, -3:]      ,						lowercase      ,						rtol=1e-1 ) )










"""simple docstring"""




__UpperCamelCase						=      frozenset(
    [
        '''prompt''',
        '''height''',
        '''width''',
        '''guidance_scale''',
        '''negative_prompt''',
        '''prompt_embeds''',
        '''negative_prompt_embeds''',
        '''cross_attention_kwargs''',
    ]
)

__UpperCamelCase						=      frozenset(['''prompt''', '''negative_prompt'''])

__UpperCamelCase						=      frozenset([])

__UpperCamelCase						=      frozenset(['''image'''])

__UpperCamelCase						=      frozenset(
    [
        '''image''',
        '''height''',
        '''width''',
        '''guidance_scale''',
    ]
)

__UpperCamelCase						=      frozenset(['''image'''])

__UpperCamelCase						=      frozenset(
    [
        '''prompt''',
        '''image''',
        '''height''',
        '''width''',
        '''guidance_scale''',
        '''negative_prompt''',
        '''prompt_embeds''',
        '''negative_prompt_embeds''',
    ]
)

__UpperCamelCase						=      frozenset(['''prompt''', '''image''', '''negative_prompt'''])

__UpperCamelCase						=      frozenset(
    [
        # Text guided image variation with an image mask
        '''prompt''',
        '''image''',
        '''mask_image''',
        '''height''',
        '''width''',
        '''guidance_scale''',
        '''negative_prompt''',
        '''prompt_embeds''',
        '''negative_prompt_embeds''',
    ]
)

__UpperCamelCase						=      frozenset(['''prompt''', '''image''', '''mask_image''', '''negative_prompt'''])

__UpperCamelCase						=      frozenset(
    [
        # image variation with an image mask
        '''image''',
        '''mask_image''',
        '''height''',
        '''width''',
        '''guidance_scale''',
    ]
)

__UpperCamelCase						=      frozenset(['''image''', '''mask_image'''])

__UpperCamelCase						=      frozenset(
    [
        '''example_image''',
        '''image''',
        '''mask_image''',
        '''height''',
        '''width''',
        '''guidance_scale''',
    ]
)

__UpperCamelCase						=      frozenset(['''example_image''', '''image''', '''mask_image'''])

__UpperCamelCase						=      frozenset(['''class_labels'''])

__UpperCamelCase						=      frozenset(['''class_labels'''])

__UpperCamelCase						=      frozenset(['''batch_size'''])

__UpperCamelCase						=      frozenset([])

__UpperCamelCase						=      frozenset(['''batch_size'''])

__UpperCamelCase						=      frozenset([])

__UpperCamelCase						=      frozenset(
    [
        '''prompt''',
        '''audio_length_in_s''',
        '''guidance_scale''',
        '''negative_prompt''',
        '''prompt_embeds''',
        '''negative_prompt_embeds''',
        '''cross_attention_kwargs''',
    ]
)

__UpperCamelCase						=      frozenset(['''prompt''', '''negative_prompt'''])
__UpperCamelCase						=      frozenset(['''input_tokens'''])

__UpperCamelCase						=      frozenset(['''input_tokens'''])


