


import inspect
import tempfile
import unittest

from huggingface_hub import hf_hub_download

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

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


a					  =     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]							,					UpperCamelCase__       :   Dict							,					UpperCamelCase__       :   Optional[int]=16							,					UpperCamelCase__       :   Optional[Any]=13							,					UpperCamelCase__       :   Tuple=7							,					UpperCamelCase__       :   List[str]=14							,					UpperCamelCase__       :   Union[str, Any]=10							,					UpperCamelCase__       :   List[Any]=19							,					UpperCamelCase__       :   Optional[int]=5							,					UpperCamelCase__       :   Optional[Any]=4							,					UpperCamelCase__       :   Any=True							,					UpperCamelCase__       :   int=16							,					UpperCamelCase__       :   Union[str, Any]=2							,					UpperCamelCase__       :   Any=4							,					UpperCamelCase__       :   Optional[Any]=4							,					UpperCamelCase__       :   Union[str, Any]="gelu"							,					UpperCamelCase__       :   int=0.1							,					UpperCamelCase__       :   str=0.1							,					UpperCamelCase__       :   str=[1, 2, 3, 4, 5]							,					UpperCamelCase__       :   str=25							,					UpperCamelCase__       :   Dict=5							,					):






										'''simple docstring'''






										lowercase_  							=			d_model
										lowercase_  							=			parent
										lowercase_  							=			batch_size
										lowercase_  							=			prediction_length
										lowercase_  							=			context_length
										lowercase_  							=			cardinality
										lowercase_  							=			num_time_features
										lowercase_  							=			lags_sequence
										lowercase_  							=			embedding_dimension
										lowercase_  							=			is_training
										lowercase_  							=			hidden_size
										lowercase_  							=			num_hidden_layers
										lowercase_  							=			num_attention_heads
										lowercase_  							=			intermediate_size
										lowercase_  							=			hidden_act
										lowercase_  							=			hidden_dropout_prob
										lowercase_  							=			attention_probs_dropout_prob

										lowercase_  							=			context_length
										lowercase_  							=			prediction_length + label_length
										lowercase_  							=			label_length

										lowercase_  							=			moving_average
										lowercase_  							=			autocorrelation_factor

					def   UpperCAmelCase__       (     self       :   str       ):






										'''simple docstring'''






										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       :   Tuple							,					UpperCamelCase__       :   Optional[Any]       ):






										'''simple docstring'''






										lowercase_  							=			config.context_length + max(config.lags_sequence       )

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

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

										lowercase_  							=			{
										    """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       :   Optional[Any]       ):






										'''simple docstring'''






										lowercase_  							=			self.get_config()
										lowercase_  							=			self.prepare_autoformer_inputs_dict(UpperCamelCase__       )
										return config, inputs_dict

					def   UpperCAmelCase__       (     self       :   Optional[Any]       ):






										'''simple docstring'''






										lowercase_							,					lowercase_  							=			self.prepare_config_and_inputs()
										return config, inputs_dict







					def   UpperCAmelCase__       (     self       :   Tuple							,					UpperCamelCase__       :   Tuple							,					UpperCamelCase__       :   int       ):






										'''simple docstring'''






										lowercase_  							=			AutoformerModel(config=UpperCamelCase__       ).to(UpperCamelCase__       ).eval()
										lowercase_  							=			model(**UpperCamelCase__       )

										lowercase_  							=			outputs.encoder_last_hidden_state
										lowercase_  							=			outputs.last_hidden_state

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

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

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

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

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

										lowercase_  							=			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        UpperCamelCase__   (	__magic_name__							,				__magic_name__							,				unittest.TestCase   ):
					__SCREAMING_SNAKE_CASE       :  str						=   (AutoformerModel, AutoformerForPrediction) if is_torch_available() else ()
					__SCREAMING_SNAKE_CASE       :  str						=   (AutoformerForPrediction,) if is_torch_available() else ()
					__SCREAMING_SNAKE_CASE       :  Tuple						=   {'feature-extraction': AutoformerModel} if is_torch_available() else {}
					__SCREAMING_SNAKE_CASE       :  str						=   False
					__SCREAMING_SNAKE_CASE       :  int						=   False
					__SCREAMING_SNAKE_CASE       :  Tuple						=   False
					__SCREAMING_SNAKE_CASE       :  Optional[Any]						=   False
					__SCREAMING_SNAKE_CASE       :  Tuple						=   False
					__SCREAMING_SNAKE_CASE       :  Optional[Any]						=   False

					def   UpperCAmelCase__       (     self       :   Tuple       ):






										'''simple docstring'''






										lowercase_  							=			AutoformerModelTester(self       )
										lowercase_  							=			ConfigTester(self							,					config_class=UpperCamelCase__							,					has_text_modality=UpperCamelCase__       )

					def   UpperCAmelCase__       (     self       :   Dict       ):






										'''simple docstring'''






										self.config_tester.run_common_tests()

					def   UpperCAmelCase__       (     self       :   Optional[int]       ):






										'''simple docstring'''






										lowercase_							,					lowercase_  							=			self.model_tester.prepare_config_and_inputs()
										for model_class in self.all_model_classes:
															lowercase_  							=			model_class(UpperCamelCase__       )

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

					def   UpperCAmelCase__       (     self       :   List[Any]       ):






										'''simple docstring'''






										lowercase_  							=			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       :   Tuple       ):






										'''simple docstring'''






										pass

					def   UpperCAmelCase__       (     self       :   List[Any]       ):






										'''simple docstring'''






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

					def   UpperCAmelCase__       (     self       :   Optional[Any]       ):






										'''simple docstring'''






										lowercase_							,					lowercase_  							=			self.model_tester.prepare_config_and_inputs_for_common()

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

															lowercase_  							=			[
															    """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       :   Optional[int]       ):






										'''simple docstring'''






										lowercase_							,					lowercase_  							=			self.model_tester.prepare_config_and_inputs_for_common()
										lowercase_  							=			True

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

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

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

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

										lowercase_  							=			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       :   List[Any]       ):






										'''simple docstring'''






										super().test_retain_grad_hidden_states_attentions()



def 			UpperCAmelCase_				(      UpperCAmelCase__="train-batch.pt"  ):
					lowercase_  							=			hf_hub_download(repo_id="""hf-internal-testing/tourism-monthly-batch"""						,							filename=UpperCAmelCase__						,							repo_type="""dataset"""  )
					lowercase_  							=			torch.load(UpperCAmelCase__						,							map_location=UpperCAmelCase__  )
					return batch







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

					def   UpperCAmelCase__       (     self       :   Dict       ):






										'''simple docstring'''






										lowercase_  							=			AutoformerModel.from_pretrained("""huggingface/autoformer-tourism-monthly"""       ).to(UpperCamelCase__       )
										lowercase_  							=			prepare_batch()

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

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

										lowercase_  							=			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       :   Any       ):






										'''simple docstring'''






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

										lowercase_  							=			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       :   List[str]       ):






										'''simple docstring'''






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

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






# Copyright 2022 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import os
import platform

import numpy as np
import psutil
import torch

from accelerate import __version__ as version
from accelerate.commands.config import default_config_file, load_config_from_file

from ..utils import is_npu_available, is_xpu_available


def 			UpperCAmelCase_				(      UpperCAmelCase__=None  ):
					if subparsers is not None:
										lowercase_  							=			subparsers.add_parser("""env"""  )
					else:
										lowercase_  							=			argparse.ArgumentParser("""Accelerate env command"""  )

					parser.add_argument(
					    """--config_file"""						,							default=UpperCAmelCase__						,							help="""The config file to use for the default values in the launching script."""  )

					if subparsers is not None:
										parser.set_defaults(func=UpperCAmelCase__  )
					return parser


def 			UpperCAmelCase_				(      UpperCAmelCase__  ):
					lowercase_  							=			torch.__version__
					lowercase_  							=			torch.cuda.is_available()
					lowercase_  							=			is_xpu_available()
					lowercase_  							=			is_npu_available()

					lowercase_  							=			"""Not found"""
					# Get the default from the config file.
					if args.config_file is not None or os.path.isfile(UpperCAmelCase__  ):
										lowercase_  							=			load_config_from_file(args.config_file  ).to_dict()

					lowercase_  							=			{
					    """`Accelerate` version""": version,
					    """Platform""": platform.platform(),
					    """Python version""": platform.python_version(),
					    """Numpy version""": np.__version__,
					    """PyTorch version (GPU?)""": F'''{pt_version} ({pt_cuda_available})''',
					    """PyTorch XPU available""": str(UpperCAmelCase__  ),
					    """PyTorch NPU available""": str(UpperCAmelCase__  ),
					    """System RAM""": F'''{psutil.virtual_memory().total / 1_0_2_4 ** 3:.2f} GB''',
					}
					if pt_cuda_available:
										lowercase_  							=			torch.cuda.get_device_name()

					print("""\nCopy-and-paste the text below in your GitHub issue\n"""  )
					print("""\n""".join([F'''- {prop}: {val}''' for prop, val in info.items()]  )  )

					print("""- `Accelerate` default config:""" if args.config_file is None else """- `Accelerate` config passed:"""  )
					lowercase_  							=			(
					    """\n""".join([F'''\t- {prop}: {val}''' for prop, val in accelerate_config.items()]  )
					    if isinstance(UpperCAmelCase__						,							UpperCAmelCase__  )
					    else F'''\t{accelerate_config}'''
					)
					print(UpperCAmelCase__  )

					lowercase_  							=			accelerate_config

					return info


def 			UpperCAmelCase_				(      ):
					lowercase_  							=			env_command_parser()
					lowercase_  							=			parser.parse_args()
					env_command(UpperCAmelCase__  )
					return 0


if __name__ == "__main__":
							raise SystemExit(main())


