


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


UpperCAmelCase    =					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 lowerCAmelCase_		:

	'''simple docstring'''

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

							snake_case_ =   context_length
							snake_case_ =   prediction_length + label_length
							snake_case_ =   label_length

							snake_case_ =   moving_average
							snake_case_ =   autocorrelation_factor



	def 		UpperCamelCase__  (			self      ):
							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	,   _UpperCAmelCase      ):
							snake_case_ =   config.context_length + max(config.lags_sequence      )

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

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

							snake_case_ =   {
							    '''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      ):
							snake_case_ =   self.get_config()
							snake_case_ =   self.prepare_autoformer_inputs_dict(_UpperCAmelCase      )
							return config, inputs_dict



	def 		UpperCamelCase__  (			self      ):
							snake_case_ ,					snake_case_ =   self.prepare_config_and_inputs()
							return config, inputs_dict





	def 		UpperCamelCase__  (			self	,   _UpperCAmelCase	,   _UpperCAmelCase      ):
							snake_case_ =   AutoformerModel(config=_UpperCAmelCase      ).to(_UpperCAmelCase      ).eval()
							snake_case_ =   model(**_UpperCAmelCase      )

							snake_case_ =   outputs.encoder_last_hidden_state
							snake_case_ =   outputs.last_hidden_state

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

							snake_case_ ,					snake_case_ ,					snake_case_ ,					snake_case_ ,					snake_case_ =   model.create_network_inputs(**_UpperCAmelCase      )
							snake_case_ ,					snake_case_ =   model.decomposition_layer(transformer_inputs[:, : config.context_length, ...]      )

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

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

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

							snake_case_ =   decoder(
							    trend=_UpperCAmelCase	,   inputs_embeds=_UpperCAmelCase	,   encoder_hidden_states=_UpperCAmelCase	,   )[0]

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







@require_torch
class lowerCAmelCase_		(				lowerCamelCase__      ,		lowerCamelCase__      ,		unittest.TestCase ):

	'''simple docstring'''

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



	def 		UpperCamelCase__  (			self      ):
							snake_case_ =   AutoformerModelTester(self      )
							snake_case_ =   ConfigTester(self	,   config_class=_UpperCAmelCase	,   has_text_modality=_UpperCAmelCase      )



	def 		UpperCamelCase__  (			self      ):
							self.config_tester.run_common_tests()



	def 		UpperCamelCase__  (			self      ):
							snake_case_ ,					snake_case_ =   self.model_tester.prepare_config_and_inputs()
							for model_class in self.all_model_classes:
													snake_case_ =   model_class(_UpperCAmelCase      )

													with tempfile.TemporaryDirectory() as tmpdirname:
																			model.save_pretrained(_UpperCAmelCase      )
																			snake_case_ ,					snake_case_ =   model_class.from_pretrained(_UpperCAmelCase	,   output_loading_info=_UpperCAmelCase      )
													self.assertEqual(info['''missing_keys''']	,   []      )



	def 		UpperCamelCase__  (			self      ):
							snake_case_ =   self.model_tester.prepare_config_and_inputs_for_common()
							self.model_tester.check_encoder_decoder_model_standalone(*_UpperCAmelCase      )



	@unittest.skip(reason='''Model has no tokens embeddings'''      )
	def 		UpperCamelCase__  (			self      ):
							pass



	def 		UpperCamelCase__  (			self      ):
							snake_case_ =   inspect.signature(getattr(_UpperCAmelCase	,   '''forward'''      )      )
							# The main input is the name of the argument after `self`
							snake_case_ =   list(model_signature.parameters.keys()      )[1]
							self.assertEqual(AutoformerModel.main_input_name	,   _UpperCAmelCase      )



	def 		UpperCamelCase__  (			self      ):
							snake_case_ ,					snake_case_ =   self.model_tester.prepare_config_and_inputs_for_common()

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

													snake_case_ =   [
													    '''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(_UpperCAmelCase      )]	,   _UpperCAmelCase      )



	def 		UpperCamelCase__  (			self      ):
							snake_case_ ,					snake_case_ =   self.model_tester.prepare_config_and_inputs_for_common()
							snake_case_ =   True

							snake_case_ =   getattr(self.model_tester	,   '''seq_length'''	,   _UpperCAmelCase      )
							snake_case_ =   getattr(self.model_tester	,   '''decoder_seq_length'''	,   _UpperCAmelCase      )
							snake_case_ =   getattr(self.model_tester	,   '''encoder_seq_length'''	,   _UpperCAmelCase      )
							snake_case_ =   getattr(self.model_tester	,   '''d_model'''	,   _UpperCAmelCase      )
							snake_case_ =   getattr(self.model_tester	,   '''num_attention_heads'''	,   _UpperCAmelCase      )
							snake_case_ =   d_model // num_attention_heads

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

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

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

													snake_case_ =   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(_UpperCAmelCase	,   _UpperCAmelCase      )

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

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

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

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






def  __lowerCAmelCase					(SCREAMING_SNAKE_CASE="train-batch.pt"      )->						int:




						"""simple docstring"""

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







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

	'''simple docstring'''

	def 		UpperCamelCase__  (			self      ):
							snake_case_ =   AutoformerModel.from_pretrained('''huggingface/autoformer-tourism-monthly'''      ).to(_UpperCAmelCase      )
							snake_case_ =   prepare_batch()

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

							snake_case_ =   torch.Size(
							    (64, model.config.prediction_length + model.config.label_length, model.config.feature_size)      )
							self.assertEqual(output.shape	,   _UpperCAmelCase      )

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



	def 		UpperCamelCase__  (			self      ):
							snake_case_ =   AutoformerForPrediction.from_pretrained('''huggingface/autoformer-tourism-monthly'''      ).to(_UpperCAmelCase      )
							snake_case_ =   prepare_batch('''val-batch.pt'''      )
							with torch.no_grad():
													snake_case_ =   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
							snake_case_ =   torch.Size((64, model.config.context_length, model.config.d_model)      )
							self.assertEqual(output.shape	,   _UpperCAmelCase      )

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





	def 		UpperCamelCase__  (			self      ):
							snake_case_ =   AutoformerForPrediction.from_pretrained('''huggingface/autoformer-tourism-monthly'''      ).to(_UpperCAmelCase      )
							snake_case_ =   prepare_batch('''val-batch.pt'''      )
							with torch.no_grad():
													snake_case_ =   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''']	,   )
							snake_case_ =   torch.Size((64, model.config.num_parallel_samples, model.config.prediction_length)      )
							self.assertEqual(outputs.sequences.shape	,   _UpperCAmelCase      )

							snake_case_ =   torch.tensor([31_30.67_63, 40_56.52_93, 70_53.07_86]	,   device=_UpperCAmelCase      )
							snake_case_ =   outputs.sequences.mean(dim=1      )
							self.assertTrue(torch.allclose(mean_prediction[0, -3:]	,   _UpperCAmelCase	,   rtol=1E-1      )      )



from dataclasses import dataclass
from typing import Optional, Tuple

import torch
from torch import nn
from transformers import RobertaPreTrainedModel, XLMRobertaConfig, XLMRobertaModel
from transformers.utils import ModelOutput







@dataclass
class lowerCAmelCase_		(				lowerCamelCase__ ):

 '''simple docstring'''

 __snake_case				=						None
 __snake_case				=						None
 __snake_case				=						None
 __snake_case				=						None







class lowerCAmelCase_		(				lowerCamelCase__ ):

 '''simple docstring'''

 def __init__(			self	,   _UpperCAmelCase=1	,   _UpperCAmelCase=0	,   _UpperCAmelCase=2	,   _UpperCAmelCase=5_12	,   _UpperCAmelCase="cls"	,   _UpperCAmelCase=False	,   _UpperCAmelCase=True	,   **_UpperCAmelCase	,   ):
       super().__init__(pad_token_id=_UpperCAmelCase	,   bos_token_id=_UpperCAmelCase	,   eos_token_id=_UpperCAmelCase	,   **_UpperCAmelCase      )
       snake_case_ =   project_dim
       snake_case_ =   pooler_fn
       snake_case_ =   learn_encoder
       snake_case_ =   use_attention_mask







class lowerCAmelCase_		(				lowerCamelCase__ ):

 '''simple docstring'''

 __snake_case				=						[R"pooler", R"logit_scale"]
 __snake_case				=						[R"position_ids", R"predictions.decoder.bias"]
 __snake_case				=						"roberta"
 __snake_case				=						RobertaSeriesConfig



 def __init__(			self	,   _UpperCAmelCase      ):
       super().__init__(_UpperCAmelCase      )
       snake_case_ =   XLMRobertaModel(_UpperCAmelCase      )
       snake_case_ =   nn.Linear(config.hidden_size	,   config.project_dim      )
       snake_case_ =   getattr(_UpperCAmelCase	,   '''has_pre_transformation'''	,   _UpperCAmelCase      )
       if self.has_pre_transformation:
             snake_case_ =   nn.Linear(config.hidden_size	,   config.project_dim      )
             snake_case_ =   nn.LayerNorm(config.hidden_size	,   eps=config.layer_norm_eps      )
       self.post_init()





 def 		UpperCamelCase__  (			self	,   _UpperCAmelCase = None	,   _UpperCAmelCase = None	,   _UpperCAmelCase = None	,   _UpperCAmelCase = None	,   _UpperCAmelCase = None	,   _UpperCAmelCase = None	,   _UpperCAmelCase = None	,   _UpperCAmelCase = None	,   _UpperCAmelCase = None	,   _UpperCAmelCase = None	,   _UpperCAmelCase = None	,   ):
       snake_case_ =   return_dict if return_dict is not None else self.config.use_return_dict

       snake_case_ =   self.base_model(
           input_ids=_UpperCAmelCase	,   attention_mask=_UpperCAmelCase	,   token_type_ids=_UpperCAmelCase	,   position_ids=_UpperCAmelCase	,   head_mask=_UpperCAmelCase	,   inputs_embeds=_UpperCAmelCase	,   encoder_hidden_states=_UpperCAmelCase	,   encoder_attention_mask=_UpperCAmelCase	,   output_attentions=_UpperCAmelCase	,   output_hidden_states=True if self.has_pre_transformation else output_hidden_states	,   return_dict=_UpperCAmelCase	,   )

       if self.has_pre_transformation:
             snake_case_ =   outputs['''hidden_states'''][-2]
             snake_case_ =   self.pre_LN(_UpperCAmelCase      )
             snake_case_ =   self.transformation_pre(_UpperCAmelCase      )

             return TransformationModelOutput(
                 projection_state=_UpperCAmelCase	,   last_hidden_state=outputs.last_hidden_state	,   hidden_states=outputs.hidden_states	,   attentions=outputs.attentions	,   )
       else:
             snake_case_ =   self.transformation(outputs.last_hidden_state      )
             return TransformationModelOutput(
                 projection_state=_UpperCAmelCase	,   last_hidden_state=outputs.last_hidden_state	,   hidden_states=outputs.hidden_states	,   attentions=outputs.attentions	,   )