





"""simple docstring"""
import logging
import os
import sys
from dataclasses import dataclass, field
from importlib import import_module
from typing import Dict, List, Optional, Tuple

import numpy as np
from seqeval.metrics import accuracy_score, fa_score, precision_score, recall_score
from torch import nn
from utils_ner import Split, TokenClassificationDataset, TokenClassificationTask

import transformers
from transformers import (
    AutoConfig,
    AutoModelForTokenClassification,
    AutoTokenizer,
    DataCollatorWithPadding,
    EvalPrediction,
    HfArgumentParser,
    Trainer,
    TrainingArguments,
    set_seed,
)
from transformers.trainer_utils import is_main_process


_snake_case     :		Tuple       						= logging.getLogger(__name__)






@dataclass
class 					_UpperCAmelCase   :
					UpperCamelCase      =							field(
					    metadata={'''help''': '''Path to pretrained model or model identifier from huggingface.co/models'''} )
					UpperCamelCase      =							field(
					    default=lowercase_     ,							metadata={'''help''': '''Pretrained config name or path if not the same as model_name'''} )
					UpperCamelCase      =							field(
					    default='''NER'''     ,							metadata={'''help''': '''Task type to fine tune in training (e.g. NER, POS, etc)'''} )
					UpperCamelCase      =							field(
					    default=lowercase_     ,							metadata={'''help''': '''Pretrained tokenizer name or path if not the same as model_name'''} )
					UpperCamelCase      =							field(default=lowercase_     ,							metadata={'''help''': '''Set this flag to use fast tokenization.'''} )
					# If you want to tweak more attributes on your tokenizer, you should do it in a distinct script,
					# or just modify its tokenizer_config.json.
					UpperCamelCase      =							field(
					    default=lowercase_     ,							metadata={'''help''': '''Where do you want to store the pretrained models downloaded from huggingface.co'''}     ,							)







@dataclass
class 					_UpperCAmelCase   :
					UpperCamelCase      =							field(
					    metadata={'''help''': '''The input data dir. Should contain the .txt files for a CoNLL-2003-formatted task.'''} )
					UpperCamelCase      =							field(
					    default=lowercase_     ,							metadata={'''help''': '''Path to a file containing all labels. If not specified, CoNLL-2003 labels are used.'''}     ,							)
					UpperCamelCase      =							field(
					    default=1_2_8     ,							metadata={
					        '''help''': (
					            '''The maximum total input sequence length after tokenization. Sequences longer '''
					            '''than this will be truncated, sequences shorter will be padded.'''
					        )
					    }     ,							)
					UpperCamelCase      =							field(
					    default=lowercase_     ,							metadata={'''help''': '''Overwrite the cached training and evaluation sets'''} )







def 					A__				(			):
							A		=  HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)       )
							if len(sys.argv       ) == 2 and sys.argv[1].endswith(".json"       ):
														# If we pass only one argument to the script and it's the path to a json file,
														# let's parse it to get our arguments.
														A,	A,	A		=  parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]       )       )
							else:
														A,	A,	A		=  parser.parse_args_into_dataclasses()

							if (
							    os.path.exists(training_args.output_dir       )
							    and os.listdir(training_args.output_dir       )
							    and training_args.do_train
							    and not training_args.overwrite_output_dir
							):
														raise ValueError(
														    F"Output directory ({training_args.output_dir}) already exists and is not empty. Use"
														    " --overwrite_output_dir to overcome."       )

							A		=  import_module("tasks"       )
							try:
														A		=  getattr(A__					,  model_args.task_type       )
														A		=  token_classification_task_clazz()
							except AttributeError:
														raise ValueError(
														    F"Task {model_args.task_type} needs to be defined as a TokenClassificationTask subclass in {module}. "
														    F"Available tasks classes are: {TokenClassificationTask.__subclasses__()}"       )

							# Setup logging
							logging.basicConfig(
							    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s"					,  datefmt="%m/%d/%Y %H:%M:%S"					,  level=logging.INFO if training_args.local_rank in [-1, 0] else logging.WARN					,  )
							logger.warning(
							    "Process rank: %s, device: %s, n_gpu: %s, distributed training: %s, 16-bits training: %s"					,  training_args.local_rank					,  training_args.device					,  training_args.n_gpu					,  bool(training_args.local_rank != -1       )					,  training_args.fpaa					,  )
							# Set the verbosity to info of the Transformers logger (on main process only):
							if is_main_process(training_args.local_rank       ):
														transformers.utils.logging.set_verbosity_info()
														transformers.utils.logging.enable_default_handler()
														transformers.utils.logging.enable_explicit_format()
							logger.info("Training/evaluation parameters %s"					,  A__       )

							# Set seed
							set_seed(training_args.seed       )

							# Prepare CONLL-2003 task
							A		=  token_classification_task.get_labels(data_args.labels       )
							A		=  dict(enumerate(A__       )       )
							A		=  len(A__       )

							# Load pretrained model and tokenizer
							#
							# Distributed training:
							# The .from_pretrained methods guarantee that only one local process can concurrently
							# download model & vocab.

							A		=  AutoConfig.from_pretrained(
							    model_args.config_name if model_args.config_name else model_args.model_name_or_path					,  num_labels=A__					,  idalabel=A__					,  labelaid={label: i for i, label in enumerate(A__       )}					,  cache_dir=model_args.cache_dir					,  )
							A		=  AutoTokenizer.from_pretrained(
							    model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path					,  cache_dir=model_args.cache_dir					,  use_fast=model_args.use_fast					,  )
							A		=  AutoModelForTokenClassification.from_pretrained(
							    model_args.model_name_or_path					,  from_tf=bool(".ckpt" in model_args.model_name_or_path       )					,  config=A__					,  cache_dir=model_args.cache_dir					,  )

							# Get datasets
							A		=  (
							    TokenClassificationDataset(
							        token_classification_task=A__					,  data_dir=data_args.data_dir					,  tokenizer=A__					,  labels=A__					,  model_type=config.model_type					,  max_seq_length=data_args.max_seq_length					,  overwrite_cache=data_args.overwrite_cache					,  mode=Split.train					,  )
							    if training_args.do_train
							    else None
							)
							A		=  (
							    TokenClassificationDataset(
							        token_classification_task=A__					,  data_dir=data_args.data_dir					,  tokenizer=A__					,  labels=A__					,  model_type=config.model_type					,  max_seq_length=data_args.max_seq_length					,  overwrite_cache=data_args.overwrite_cache					,  mode=Split.dev					,  )
							    if training_args.do_eval
							    else None
							)

							def align_predictions(UpperCamelCase					,  UpperCamelCase       ) -> Tuple[List[int], List[int]]:
														A		=  np.argmax(A__					,  axis=2       )

														A,	A		=  preds.shape

														A		=  [[] for _ in range(A__       )]
														A		=  [[] for _ in range(A__       )]

														for i in range(A__       ):
																					for j in range(A__       ):
																												if label_ids[i, j] != nn.CrossEntropyLoss().ignore_index:
																																			out_label_list[i].append(label_map[label_ids[i][j]]       )
																																			preds_list[i].append(label_map[preds[i][j]]       )

														return preds_list, out_label_list

							def compute_metrics(UpperCamelCase       ) -> Dict:
														A,	A		=  align_predictions(p.predictions					,  p.label_ids       )
														return {
														    "accuracy_score": accuracy_score(A__					,  A__       ),
														    "precision": precision_score(A__					,  A__       ),
														    "recall": recall_score(A__					,  A__       ),
														    "f1": fa_score(A__					,  A__       ),
														}

							# Data collator
							A		=  DataCollatorWithPadding(A__					,  pad_to_multiple_of=8       ) if training_args.fpaa else None

							# Initialize our Trainer
							A		=  Trainer(
							    model=A__					,  args=A__					,  train_dataset=A__					,  eval_dataset=A__					,  compute_metrics=A__					,  data_collator=A__					,  )

							# Training
							if training_args.do_train:
														trainer.train(
														    model_path=model_args.model_name_or_path if os.path.isdir(model_args.model_name_or_path       ) else None       )
														trainer.save_model()
														# For convenience, we also re-save the tokenizer to the same directory,
														# so that you can share your model easily on huggingface.co/models =)
														if trainer.is_world_process_zero():
																					tokenizer.save_pretrained(training_args.output_dir       )

    # Evaluation
							A		=  {}
							if training_args.do_eval:
														logger.info("*** Evaluate ***"       )

														A		=  trainer.evaluate()

														A		=  os.path.join(training_args.output_dir					,  "eval_results.txt"       )
														if trainer.is_world_process_zero():
																					with open(A__					,  "w"       ) as writer:
																												logger.info("***** Eval results *****"       )
																												for key, value in result.items():
																																			logger.info("  %s = %s"					,  A__					,  A__       )
																																			writer.write("%s = %s\n" % (key, value)       )

																					results.update(A__       )

    # Predict
							if training_args.do_predict:
														A		=  TokenClassificationDataset(
														    token_classification_task=A__					,  data_dir=data_args.data_dir					,  tokenizer=A__					,  labels=A__					,  model_type=config.model_type					,  max_seq_length=data_args.max_seq_length					,  overwrite_cache=data_args.overwrite_cache					,  mode=Split.test					,  )

														A,	A,	A		=  trainer.predict(A__       )
														A,	A		=  align_predictions(A__					,  A__       )

														A		=  os.path.join(training_args.output_dir					,  "test_results.txt"       )
														if trainer.is_world_process_zero():
																					with open(A__					,  "w"       ) as writer:
																												for key, value in metrics.items():
																																			logger.info("  %s = %s"					,  A__					,  A__       )
																																			writer.write("%s = %s\n" % (key, value)       )

        # Save predictions
														A		=  os.path.join(training_args.output_dir					,  "test_predictions.txt"       )
														if trainer.is_world_process_zero():
																					with open(A__					,  "w"       ) as writer:
																												with open(os.path.join(data_args.data_dir					,  "test.txt"       )					,  "r"       ) as f:
																																			token_classification_task.write_predictions_to_file(A__					,  A__					,  A__       )

							return results






def 					A__				(			UpperCamelCase       ):
							main()


if __name__ == "__main__":
		main()





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


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




  def __init__(				self							,   lowerCAmelCase							,   lowerCAmelCase=16							,   lowerCAmelCase=13							,   lowerCAmelCase=7							,   lowerCAmelCase=14							,   lowerCAmelCase=10							,   lowerCAmelCase=19							,   lowerCAmelCase=5							,   lowerCAmelCase=4							,   lowerCAmelCase=True							,   lowerCAmelCase=16							,   lowerCAmelCase=2							,   lowerCAmelCase=4							,   lowerCAmelCase=4							,   lowerCAmelCase="gelu"							,   lowerCAmelCase=0.1							,   lowerCAmelCase=0.1							,   lowerCAmelCase=[1, 2, 3, 4, 5]							,   lowerCAmelCase=25							,   lowerCAmelCase=5							,   )  -> Optional[Any]:
        '''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   A__     (				self					)  -> Optional[Any]:
        '''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   A__     (				self							,   lowerCAmelCase					)  -> Dict:
        '''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   A__     (				self					)  -> List[str]:
        '''simple docstring'''






        _lowercase	       =self.get_config()
        _lowercase	       =self.prepare_autoformer_inputs_dict(lowerCAmelCase					)
        return config, inputs_dict




  def   A__     (				self					)  -> List[str]:
        '''simple docstring'''






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




  def   A__     (				self							,   lowerCAmelCase							,   lowerCAmelCase					)  -> Tuple:
        '''simple docstring'''






        _lowercase	       =AutoformerModel(config=lowerCAmelCase					).to(lowerCAmelCase					).eval()
        _lowercase	       =model(**lowerCAmelCase					)

        _lowercase	       =outputs.encoder_last_hidden_state
        _lowercase	       =outputs.last_hidden_state

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

        _lowercase       ,       _lowercase       ,       _lowercase       ,       _lowercase       ,       _lowercase	       =model.create_network_inputs(**lowerCAmelCase					)
        _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=lowerCAmelCase					)[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(lowerCAmelCase					)
              _lowercase	       =AutoformerDecoder.from_pretrained(lowerCAmelCase					).to(lowerCAmelCase					)

        _lowercase	       =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  __lowerCAmelCase     (					SCREAMING_SNAKE_CASE			,			SCREAMING_SNAKE_CASE			,			unittest.TestCase							):
  _a									=  (AutoformerModel, AutoformerForPrediction) if is_torch_available() else ()
  _a									=  (AutoformerForPrediction,) if is_torch_available() else ()
  _a									=  {"""feature-extraction""": AutoformerModel} if is_torch_available() else {}
  _a									=  False
  _a									=  False
  _a									=  False
  _a									=  False
  _a									=  False
  _a									=  False




  def   A__     (				self					)  -> str:
        '''simple docstring'''






        _lowercase	       =AutoformerModelTester(self					)
        _lowercase	       =ConfigTester(self							,   config_class=lowerCAmelCase							,   has_text_modality=lowerCAmelCase					)




  def   A__     (				self					)  -> Optional[Any]:
        '''simple docstring'''






        self.config_tester.run_common_tests()




  def   A__     (				self					)  -> Optional[Any]:
        '''simple docstring'''






        _lowercase       ,       _lowercase	       =self.model_tester.prepare_config_and_inputs()
        for model_class in self.all_model_classes:
              _lowercase	       =model_class(lowerCAmelCase					)

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




  def   A__     (				self					)  -> List[str]:
        '''simple docstring'''






        _lowercase	       =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   A__     (				self					)  -> int:
        '''simple docstring'''






        pass




  def   A__     (				self					)  -> Dict:
        '''simple docstring'''






        _lowercase	       =inspect.signature(getattr(lowerCAmelCase							,   '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							,   lowerCAmelCase					)




  def   A__     (				self					)  -> List[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(lowerCAmelCase					)
              _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(lowerCAmelCase					)]							,   lowerCAmelCase					)




  def   A__     (				self					)  -> int:
        '''simple docstring'''






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

        _lowercase	       =getattr(self.model_tester							,   'seq_length'							,   lowerCAmelCase					)
        _lowercase	       =getattr(self.model_tester							,   'decoder_seq_length'							,   lowerCAmelCase					)
        _lowercase	       =getattr(self.model_tester							,   'encoder_seq_length'							,   lowerCAmelCase					)
        _lowercase	       =getattr(self.model_tester							,   'd_model'							,   lowerCAmelCase					)
        _lowercase	       =getattr(self.model_tester							,   'num_attention_heads'							,   lowerCAmelCase					)
        _lowercase	       =d_model // num_attention_heads

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

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

              # decoder attentions
              _lowercase	       =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
              _lowercase	       =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
        _lowercase	       =True
        _lowercase	       =True
        _lowercase	       =model_class(lowerCAmelCase					)
        model.to(lowerCAmelCase					)
        model.eval()
        with torch.no_grad():
              _lowercase	       =model(**self._prepare_for_class(lowerCAmelCase							,   lowerCAmelCase					)					)

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

        _lowercase	       =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   A__     (				self					)  -> Dict:
        '''simple docstring'''






        super().test_retain_grad_hidden_states_attentions()
def      a      (			A__	:   List[str]="train-batch.pt"     )    ->      str:





      """simple docstring"""



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







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




  def   A__     (				self					)  -> int:
        '''simple docstring'''






        _lowercase	       =AutoformerModel.from_pretrained('huggingface/autoformer-tourism-monthly'					).to(lowerCAmelCase					)
        _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							,   lowerCAmelCase					)

        _lowercase	       =torch.tensor(
            [[0.3593, -1.3398, 0.6330], [0.2279, 1.5396, -0.1792], [0.0450, 1.3225, -0.2335]]							,   device=lowerCAmelCase					)
        self.assertTrue(torch.allclose(output[0, :3, :3]							,   lowerCAmelCase							,   atol=lowerCAmelCase					)					)




  def   A__     (				self					)  -> str:
        '''simple docstring'''






        _lowercase	       =AutoformerForPrediction.from_pretrained('huggingface/autoformer-tourism-monthly'					).to(lowerCAmelCase					)
        _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							,   lowerCAmelCase					)

        _lowercase	       =torch.tensor(
            [[-0.0734, -0.9036, 0.8358], [4.7186, 2.4113, 1.9581], [1.7953, 2.3558, 1.2970]]							,   device=lowerCAmelCase					)
        self.assertTrue(torch.allclose(output[0, :3, :3]							,   lowerCAmelCase							,   atol=lowerCAmelCase					)					)




  def   A__     (				self					)  -> Optional[int]:
        '''simple docstring'''






        _lowercase	       =AutoformerForPrediction.from_pretrained('huggingface/autoformer-tourism-monthly'					).to(lowerCAmelCase					)
        _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							,   lowerCAmelCase					)

        _lowercase	       =torch.tensor([3130.6763, 4056.5293, 7053.0786]							,   device=lowerCAmelCase					)
        _lowercase	       =outputs.sequences.mean(dim=1					)
        self.assertTrue(torch.allclose(mean_prediction[0, -3:]							,   lowerCAmelCase							,   rtol=1e-1					)					)


