


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


__lowerCamelCase        =     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   a__	:

							def __init__(   self      :     Optional[int]      ,							lowerCamelCase_      :     Any      ,							lowerCamelCase_      :     Optional[int]=1_6      ,							lowerCamelCase_      :     int=1_3      ,							lowerCamelCase_      :     Union[str, Any]=7      ,							lowerCamelCase_      :     List[str]=1_4      ,							lowerCamelCase_      :     Tuple=1_0      ,							lowerCamelCase_      :     int=1_9      ,							lowerCamelCase_      :     int=5      ,							lowerCamelCase_      :     str=4      ,							lowerCamelCase_      :     Optional[int]=True      ,							lowerCamelCase_      :     str=1_6      ,							lowerCamelCase_      :     Dict=2      ,							lowerCamelCase_      :     Union[str, Any]=4      ,							lowerCamelCase_      :     str=4      ,							lowerCamelCase_      :     Union[str, Any]="gelu"      ,							lowerCamelCase_      :     List[Any]=0.1      ,							lowerCamelCase_      :     Any=0.1      ,							lowerCamelCase_      :     str=[1, 2, 3, 4, 5]      ,							lowerCamelCase_      :     List[Any]=2_5      ,							lowerCamelCase_      :     Union[str, Any]=5      ,							):
												a_    :     List[str]										=    d_model
												a_    :     Dict										=    parent
												a_    :     Optional[int]										=    batch_size
												a_    :     Union[str, Any]										=    prediction_length
												a_    :     str										=    context_length
												a_    :     Dict										=    cardinality
												a_    :     Tuple										=    num_time_features
												a_    :     Optional[Any]										=    lags_sequence
												a_    :     str										=    embedding_dimension
												a_    :     Optional[Any]										=    is_training
												a_    :     Dict										=    hidden_size
												a_    :     Union[str, Any]										=    num_hidden_layers
												a_    :     Optional[int]										=    num_attention_heads
												a_    :     Dict										=    intermediate_size
												a_    :     Dict										=    hidden_act
												a_    :     List[str]										=    hidden_dropout_prob
												a_    :     List[str]										=    attention_probs_dropout_prob

												a_    :     Optional[int]										=    context_length
												a_    :     List[Any]										=    prediction_length + label_length
												a_    :     Dict										=    label_length

												a_    :     Dict										=    moving_average
												a_    :     List[Any]										=    autocorrelation_factor

							def 						UpperCAmelCase(   self      :     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 						UpperCAmelCase(   self      :     Optional[int]      ,							lowerCamelCase_      :     Optional[Any]    ):
												a_    :     Tuple										=    config.context_length + max(config.lags_sequence    )

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

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

												a_    :     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 						UpperCAmelCase(   self      :     str    ):
												a_    :     str										=    self.get_config()
												a_    :     Optional[Any]										=    self.prepare_autoformer_inputs_dict(lowerCamelCase_    )
												return config, inputs_dict

							def 						UpperCAmelCase(   self      :     List[str]    ):
												a_    :     Optional[Any]										=    self.prepare_config_and_inputs()
												return config, inputs_dict

							def 						UpperCAmelCase(   self      :     int      ,							lowerCamelCase_      :     Union[str, Any]      ,							lowerCamelCase_      :     str    ):
												a_    :     Optional[int]										=    AutoformerModel(config=lowerCamelCase_    ).to(lowerCamelCase_    ).eval()
												a_    :     str										=    model(**lowerCamelCase_    )

												a_    :     Union[str, Any]										=    outputs.encoder_last_hidden_state
												a_    :     int										=    outputs.last_hidden_state

												with tempfile.TemporaryDirectory() as tmpdirname:
																	a_    :     str										=    model.get_encoder()
																	encoder.save_pretrained(lowerCamelCase_    )
																	a_    :     Any										=    AutoformerEncoder.from_pretrained(lowerCamelCase_    ).to(lowerCamelCase_    )

												a_    :     int										=    model.create_network_inputs(**lowerCamelCase_    )
												a_    :     Union[str, Any]										=    model.decomposition_layer(transformer_inputs[:, : config.context_length, ...]    )

												a_    :     List[str]										=    torch.cat(
												    (transformer_inputs[:, : config.context_length, ...], feature[:, : config.context_length, ...])      ,							dim=-1      ,							)
												a_    :     int										=    encoder(inputs_embeds=lowerCamelCase_    )[0]
												self.parent.assertTrue((encoder_last_hidden_state_a - encoder_last_hidden_state).abs().max().item() < 1E-3    )

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

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

												a_    :     Union[str, Any]										=    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   a__	(						lowerCAmelCase_						,       lowerCAmelCase_						,       unittest.TestCase    ):
							lowerCamelCase__:     Any							    =     (AutoformerModel, AutoformerForPrediction) if is_torch_available() else ()
							lowerCamelCase__:     List[Any]							    =     (AutoformerForPrediction,) if is_torch_available() else ()
							lowerCamelCase__:     List[Any]							    =     {"""feature-extraction""": AutoformerModel} if is_torch_available() else {}
							lowerCamelCase__:     Union[str, Any]							    =     False
							lowerCamelCase__:     str							    =     False
							lowerCamelCase__:     Optional[int]							    =     False
							lowerCamelCase__:     int							    =     False
							lowerCamelCase__:     List[Any]							    =     False
							lowerCamelCase__:     Dict							    =     False

							def 						UpperCAmelCase(   self      :     str    ):
												a_    :     str										=    AutoformerModelTester(self    )
												a_    :     List[str]										=    ConfigTester(self      ,							config_class=lowerCamelCase_      ,							has_text_modality=lowerCamelCase_    )

							def 						UpperCAmelCase(   self      :     Optional[int]    ):
												self.config_tester.run_common_tests()

							def 						UpperCAmelCase(   self      :     List[Any]    ):
												a_    :     int										=    self.model_tester.prepare_config_and_inputs()
												for model_class in self.all_model_classes:
																	a_    :     List[Any]										=    model_class(lowerCamelCase_    )

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

							def 						UpperCAmelCase(   self      :     Optional[Any]    ):
												a_    :     List[str]										=    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 						UpperCAmelCase(   self      :     Optional[Any]    ):
												pass

							def 						UpperCAmelCase(   self      :     Optional[int]    ):
												a_    :     str										=    inspect.signature(getattr(lowerCamelCase_      ,							"""forward"""    )    )
												# The main input is the name of the argument after `self`
												a_    :     Optional[Any]										=    list(model_signature.parameters.keys()    )[1]
												self.assertEqual(AutoformerModel.main_input_name      ,							lowerCamelCase_    )

							def 						UpperCAmelCase(   self      :     Any    ):
												a_    :     List[Any]										=    self.model_tester.prepare_config_and_inputs_for_common()

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

																	a_    :     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 						UpperCAmelCase(   self      :     Optional[Any]    ):
												a_    :     str										=    self.model_tester.prepare_config_and_inputs_for_common()
												a_    :     Optional[Any]										=    True

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

												for model_class in self.all_model_classes:
																	a_    :     Union[str, Any]										=    True
																	a_    :     str										=    False
																	a_    :     Optional[Any]										=    True
																	a_    :     List[str]										=    model_class(lowerCamelCase_    )
																	model.to(lowerCamelCase_    )
																	model.eval()
																	with torch.no_grad():
																						a_    :     Dict										=    model(**self._prepare_for_class(lowerCamelCase_      ,							lowerCamelCase_    )    )
																	a_    :     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"]
																	a_    :     str										=    True
																	a_    :     Optional[Any]										=    model_class(lowerCamelCase_    )
																	model.to(lowerCamelCase_    )
																	model.eval()
																	with torch.no_grad():
																						a_    :     List[str]										=    model(**self._prepare_for_class(lowerCamelCase_      ,							lowerCamelCase_    )    )
																	a_    :     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]      ,							)
																	a_    :     str										=    len(lowerCamelCase_    )

																	a_    :     List[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
																	a_    :     Optional[int]										=    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
																	a_    :     List[str]										=    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
												a_    :     List[str]										=    True
												a_    :     List[Any]										=    True
												a_    :     List[str]										=    model_class(lowerCamelCase_    )
												model.to(lowerCamelCase_    )
												model.eval()
												with torch.no_grad():
																	a_    :     Union[str, Any]										=    model(**self._prepare_for_class(lowerCamelCase_      ,							lowerCamelCase_    )    )

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

												a_    :     Any										=    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 						UpperCAmelCase(   self      :     List[Any]    ):
												super().test_retain_grad_hidden_states_attentions()



def  _a			(						__UpperCamelCase="train-batch.pt" ):
					a_    :     Optional[int]										=    hf_hub_download(repo_id="""hf-internal-testing/tourism-monthly-batch"""  ,  filename=__UpperCamelCase  ,  repo_type="""dataset""" )
					a_    :     Union[str, Any]										=    torch.load(__UpperCamelCase  ,  map_location=__UpperCamelCase )
					return batch


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

							def 						UpperCAmelCase(   self      :     int    ):
												a_    :     List[Any]										=    AutoformerModel.from_pretrained("""huggingface/autoformer-tourism-monthly"""    ).to(lowerCamelCase_    )
												a_    :     Union[str, Any]										=    prepare_batch()

												with torch.no_grad():
																	a_    :     Dict										=    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]

												a_    :     Union[str, Any]										=    torch.Size(
												    (6_4, model.config.prediction_length + model.config.label_length, model.config.feature_size)    )
												self.assertEqual(output.shape      ,							lowerCamelCase_    )

												a_    :     Dict										=    torch.tensor(
												    [[0.3_5_9_3, -1.3_3_9_8, 0.6_3_3_0], [0.2_2_7_9, 1.5_3_9_6, -0.1_7_9_2], [0.0_4_5_0, 1.3_2_2_5, -0.2_3_3_5]]      ,							device=lowerCamelCase_    )
												self.assertTrue(torch.allclose(output[0, :3, :3]      ,							lowerCamelCase_      ,							atol=lowerCamelCase_    )    )

							def 						UpperCAmelCase(   self      :     Tuple    ):
												a_    :     Optional[Any]										=    AutoformerForPrediction.from_pretrained("""huggingface/autoformer-tourism-monthly"""    ).to(lowerCamelCase_    )
												a_    :     Union[str, Any]										=    prepare_batch("""val-batch.pt"""    )
												with torch.no_grad():
																	a_    :     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"""]      ,							).encoder_last_hidden_state
												a_    :     int										=    torch.Size((6_4, model.config.context_length, model.config.d_model)    )
												self.assertEqual(output.shape      ,							lowerCamelCase_    )

												a_    :     List[Any]										=    torch.tensor(
												    [[-0.0_7_3_4, -0.9_0_3_6, 0.8_3_5_8], [4.7_1_8_6, 2.4_1_1_3, 1.9_5_8_1], [1.7_9_5_3, 2.3_5_5_8, 1.2_9_7_0]]      ,							device=lowerCamelCase_    )
												self.assertTrue(torch.allclose(output[0, :3, :3]      ,							lowerCamelCase_      ,							atol=lowerCamelCase_    )    )

							def 						UpperCAmelCase(   self      :     int    ):
												a_    :     int										=    AutoformerForPrediction.from_pretrained("""huggingface/autoformer-tourism-monthly"""    ).to(lowerCamelCase_    )
												a_    :     int										=    prepare_batch("""val-batch.pt"""    )
												with torch.no_grad():
																	a_    :     Tuple										=    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"""]      ,							)
												a_    :     int										=    torch.Size((6_4, model.config.num_parallel_samples, model.config.prediction_length)    )
												self.assertEqual(outputs.sequences.shape      ,							lowerCamelCase_    )

												a_    :     Union[str, Any]										=    torch.tensor([3_1_3_0.6_7_6_3, 4_0_5_6.5_2_9_3, 7_0_5_3.0_7_8_6]      ,							device=lowerCamelCase_    )
												a_    :     Optional[Any]										=    outputs.sequences.mean(dim=1    )
												self.assertTrue(torch.allclose(mean_prediction[0, -3:]      ,							lowerCamelCase_      ,							rtol=1E-1    )    )







from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


__lowerCamelCase        =     {
    '''configuration_pegasus_x''': ['''PEGASUS_X_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''PegasusXConfig'''],
}

try:
       if not is_torch_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       __lowerCamelCase        =     [
           '''PEGASUS_X_PRETRAINED_MODEL_ARCHIVE_LIST''',
           '''PegasusXForConditionalGeneration''',
           '''PegasusXModel''',
           '''PegasusXPreTrainedModel''',
       ]


if TYPE_CHECKING:
       from .configuration_pegasus_x import PEGASUS_X_PRETRAINED_CONFIG_ARCHIVE_MAP, PegasusXConfig

       try:
              if not is_torch_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_pegasus_x import (
                  PEGASUS_X_PRETRAINED_MODEL_ARCHIVE_LIST,
                  PegasusXForConditionalGeneration,
                  PegasusXModel,
                  PegasusXPreTrainedModel,
              )


else:
       import sys

       __lowerCamelCase        =     _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)



