





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


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

     A__    =      context_length
     A__    =      prediction_length + label_length
     A__    =      label_length

     A__    =      moving_average
     A__    =      autocorrelation_factor
 def 					UpperCAmelCase_							(    self				:List[Any]							)->							List[Any]:
     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				:Dict						,	lowercase_				:Optional[int]							)->							Optional[int]:
     A__    =      config.context_length + max(config.lags_sequence							)

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

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

     A__    =      {
         "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							)->							Optional[Any]:
     A__    =      self.get_config()
     A__    =      self.prepare_autoformer_inputs_dict(lowercase_							)
     return config, inputs_dict
 def 					UpperCAmelCase_							(    self				:Tuple							)->							Any:
     A__,     A__    =      self.prepare_config_and_inputs()
     return config, inputs_dict



 def 					UpperCAmelCase_							(    self				:str						,	lowercase_				:Dict						,	lowercase_				:int							)->							List[str]:
     A__    =      AutoformerModel(config=lowercase_							).to(lowercase_							).eval()
     A__    =      model(**lowercase_							)

     A__    =      outputs.encoder_last_hidden_state
     A__    =      outputs.last_hidden_state

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

     A__,     A__,     A__,     A__,     A__    =      model.create_network_inputs(**lowercase_							)
     A__,     A__    =      model.decomposition_layer(transformer_inputs[:, : config.context_length, ...]							)

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

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

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

     A__    =      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        UpperCAmelCase     (	UpperCamelCase__      ,	UpperCamelCase__      ,	unittest.TestCase    ):
 __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 					UpperCAmelCase_							(    self				:List[str]							)->							Optional[Any]:
     A__    =      AutoformerModelTester(self							)
     A__    =      ConfigTester(self						,	config_class=lowercase_						,	has_text_modality=lowercase_							)
 def 					UpperCAmelCase_							(    self				:Union[str, Any]							)->							Optional[int]:
     self.config_tester.run_common_tests()
 def 					UpperCAmelCase_							(    self				:Any							)->							List[Any]:
     A__,     A__    =      self.model_tester.prepare_config_and_inputs()
     for model_class in self.all_model_classes:
         A__    =      model_class(lowercase_							)

         with tempfile.TemporaryDirectory() as tmpdirname:
             model.save_pretrained(lowercase_							)
             A__,     A__    =      model_class.from_pretrained(lowercase_						,	output_loading_info=lowercase_							)
         self.assertEqual(info["missing_keys"]						,	[]							)
 def 					UpperCAmelCase_							(    self				:Optional[Any]							)->							int:
     A__    =      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 					UpperCAmelCase_							(    self				:str							)->							int:
     pass
 def 					UpperCAmelCase_							(    self				:Tuple							)->							Tuple:
     A__    =      inspect.signature(getattr(lowercase_						,	"forward"							)							)
     # The main input is the name of the argument after `self`
     A__    =      list(model_signature.parameters.keys()							)[1]
     self.assertEqual(AutoformerModel.main_input_name						,	lowercase_							)
 def 					UpperCAmelCase_							(    self				:str							)->							int:
     A__,     A__    =      self.model_tester.prepare_config_and_inputs_for_common()

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

         A__    =      [
             "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 					UpperCAmelCase_							(    self				:Optional[Any]							)->							Dict:
     A__,     A__    =      self.model_tester.prepare_config_and_inputs_for_common()
     A__    =      True

     A__    =      getattr(self.model_tester						,	"seq_length"						,	lowercase_							)
     A__    =      getattr(self.model_tester						,	"decoder_seq_length"						,	lowercase_							)
     A__    =      getattr(self.model_tester						,	"encoder_seq_length"						,	lowercase_							)
     A__    =      getattr(self.model_tester						,	"d_model"						,	lowercase_							)
     A__    =      getattr(self.model_tester						,	"num_attention_heads"						,	lowercase_							)
     A__    =      d_model // num_attention_heads

     for model_class in self.all_model_classes:
         A__    =      True
         A__    =      False
         A__    =      True
         A__    =      model_class(lowercase_							)
         model.to(lowercase_							)
         model.eval()
         with torch.no_grad():
             A__    =      model(**self._prepare_for_class(lowercase_						,	lowercase_							)							)
         A__    =      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"]
         A__    =      True
         A__    =      model_class(lowercase_							)
         model.to(lowercase_							)
         model.eval()
         with torch.no_grad():
             A__    =      model(**self._prepare_for_class(lowercase_						,	lowercase_							)							)
         A__    =      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]						,	)
         A__    =      len(lowercase_							)

         A__    =      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
         A__    =      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
         A__    =      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
     A__    =      True
     A__    =      True
     A__    =      model_class(lowercase_							)
     model.to(lowercase_							)
     model.eval()
     with torch.no_grad():
         A__    =      model(**self._prepare_for_class(lowercase_						,	lowercase_							)							)

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

     A__    =      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 					UpperCAmelCase_							(    self				:str							)->							Optional[Any]:
     super().test_retain_grad_hidden_states_attentions()





def 						UpperCamelCase    (							_lowerCamelCase : int="train-batch.pt"	):
    A__    =      hf_hub_download(repo_id="hf-internal-testing/tourism-monthly-batch"						,	filename=_lowerCamelCase						,	repo_type="dataset"	)
    A__    =      torch.load(_lowerCamelCase						,	map_location=_lowerCamelCase	)
    return batch






@require_torch
@slow
class        UpperCAmelCase     (	unittest.TestCase    ):
 def 					UpperCAmelCase_							(    self				:Tuple							)->							Optional[int]:
     A__    =      AutoformerModel.from_pretrained("huggingface/autoformer-tourism-monthly"							).to(lowercase_							)
     A__    =      prepare_batch()

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

     A__    =      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=lowercase_							)
     self.assertTrue(torch.allclose(output[0, :3, :3]						,	lowercase_						,	atol=lowercase_							)							)
 def 					UpperCAmelCase_							(    self				:str							)->							List[str]:
     A__    =      AutoformerForPrediction.from_pretrained("huggingface/autoformer-tourism-monthly"							).to(lowercase_							)
     A__    =      prepare_batch("val-batch.pt"							)
     with torch.no_grad():
         A__    =      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__    =      torch.Size((64, model.config.context_length, model.config.d_model)							)
     self.assertEqual(output.shape						,	lowercase_							)

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



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

     A__    =      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=lowercase_							)
     A__    =      outputs.sequences.mean(dim=1							)
     self.assertTrue(torch.allclose(mean_prediction[0, -3:]						,	lowercase_						,	rtol=1E-1							)							)










'''simple docstring'''

from .imports import is_tqdm_available


if is_tqdm_available():
  from tqdm.auto import tqdm as _tqdm

from ..state import PartialState





def 						UpperCamelCase    (							_lowerCamelCase : bool = True						,	*_lowerCamelCase : Optional[int]						,	**_lowerCamelCase : Optional[Any]	):
    if not is_tqdm_available():
        raise ImportError("Accelerate's `tqdm` module requires `tqdm` to be installed. Please run `pip install tqdm`."	)
    A__    =      False
    if main_process_only:
        A__    =      PartialState().local_process_index == 0
    return _tqdm(*_lowerCamelCase						,	**_lowerCamelCase						,	disable=_lowerCamelCase	)



