

import os

import jsonlines
import numpy as np
from tqdm import tqdm


lowercase_									=   2_0_4_8
lowercase_									=   4_0_9_6
lowercase_									=   4_2
lowercase_									=   os.environ.pop('PROCESS_TRAIN', 'false')
lowercase_									=   {'null': 0, 'short': 1, 'long': 2, 'yes': 3, 'no': 4}
def      a      (			A__	:   Union[str, Any]     )    ->      int:





      """simple docstring"""



      def choose_first(A__	:   Optional[Any]      ,   A__	:   Tuple=False     ):
            assert isinstance(A__      ,   A__     )
            if len(A__     ) == 1:
                  _lowercase	       =answer[0]
                  return {k: [answer[k]] for k in answer} if is_long_answer else answer
            for a in answer:
                  if is_long_answer:
                        _lowercase	       ={k: [a[k]] for k in a}
                  if len(a['start_token']     ) > 0:
                        break
            return a

      _lowercase	       ={'id': example['id']}
      _lowercase	       =example['annotations']
      _lowercase	       =annotation['yes_no_answer']
      if 0 in yes_no_answer or 1 in yes_no_answer:
            _lowercase	       =['yes'] if 1 in yes_no_answer else ['no']
            _lowercase	       =_lowercase	       =[]
            _lowercase	       =_lowercase	       =[]
            _lowercase	       =['<cls>']
      else:
            _lowercase	       =['short']
            _lowercase	       =choose_first(annotation['short_answers']     )
            if len(out['start_token']     ) == 0:
                  # answer will be long if short is not available
                  _lowercase	       =['long']
                  _lowercase	       =choose_first(annotation['long_answer']      ,   is_long_answer=A__     )
                  _lowercase	       =[]
            answer.update(A__     )

      # disregard some samples
      if len(answer['start_token']     ) > 1 or answer["start_token"] == answer["end_token"]:
            _lowercase	       =True
      else:
            _lowercase	       =False

      _lowercase	       =['start_token', 'end_token', 'start_byte', 'end_byte', 'text']
      if not all(isinstance(answer[k]      ,   A__     ) for k in cols     ):
            raise ValueError('Issue in ID'      ,   example['id']     )

      return answer
def      a      (			A__	:   Union[str, Any]      ,   A__	:   Optional[int]=False     )    ->      Optional[int]:





      """simple docstring"""



      _lowercase	       =_get_single_answer(A__     )
      # bytes are of no use
      del answer["start_byte"]
      del answer["end_byte"]

      # handle yes_no answers explicitly
      if answer["category"][0] in ["yes", "no"]:  # category is list with one element
            _lowercase	       =example['document']['tokens']
            _lowercase	       =[]
            for i in range(len(doc['token']     )     ):
                  if not doc["is_html"][i]:
                        context.append(doc['token'][i]     )
            return {
                "context": " ".join(A__     ),
                "answer": {
                    "start_token": -100,  # ignore index in cross-entropy
                    "end_token": -100,  # ignore index in cross-entropy
                    "category": answer["category"],
                    "span": answer["category"],  # extra
                },
            }

      # later, help in removing all no answers
      if answer["start_token"] == [-1]:
            return {
                "context": "None",
                "answer": {
                    "start_token": -1,
                    "end_token": -1,
                    "category": "null",
                    "span": "None",  # extra
                },
            }

      # handling normal samples

      _lowercase	       =['start_token', 'end_token']
      answer.update({k: answer[k][0] if len(answer[k]     ) > 0 else answer[k] for k in cols}     )  # e.g. [10] == 10

      _lowercase	       =example['document']['tokens']
      _lowercase	       =answer['start_token']
      _lowercase	       =answer['end_token']

      _lowercase	       =[]
      for i in range(len(doc['token']     )     ):
            if not doc["is_html"][i]:
                  context.append(doc['token'][i]     )
            else:
                  if answer["start_token"] > i:
                        start_token -= 1
                  if answer["end_token"] > i:
                        end_token -= 1
      _lowercase	       =' '.join(context[start_token:end_token]     )

      # checking above code
      if assertion:
            _lowercase	       =doc['is_html'][answer['start_token'] : answer['end_token']]
            _lowercase	       =doc['token'][answer['start_token'] : answer['end_token']]
            _lowercase	       =' '.join([old[i] for i in range(len(A__     )     ) if not is_html[i]]     )
            if new != old:
                  print('ID:'      ,   example['id']     )
                  print('New:'      ,   A__      ,   end='\n'     )
                  print('Old:'      ,   A__      ,   end='\n\n'     )

      return {
          "context": " ".join(A__     ),
          "answer": {
              "start_token": start_token,
              "end_token": end_token - 1,  # this makes it inclusive
              "category": answer["category"],  # either long or short
              "span": new,  # extra
          },
      }
def      a      (			A__	:   Union[str, Any]      ,   A__	:   Tuple      ,   A__	:   Any=2048      ,   A__	:   Dict=4096      ,   A__	:   Dict=True     )    ->      Any:





      """simple docstring"""



      _lowercase	       =get_context_and_ans(A__      ,   assertion=A__     )
      _lowercase	       =out['answer']

      # later, removing these samples
      if answer["start_token"] == -1:
            return {
                "example_id": example["id"],
                "input_ids": [[-1]],
                "labels": {
                    "start_token": [-1],
                    "end_token": [-1],
                    "category": ["null"],
                },
            }

      _lowercase	       =tokenizer(example['question']['text']      ,   out['context']     ).input_ids
      _lowercase	       =input_ids.index(tokenizer.sep_token_id     ) + 1

      # return yes/no
      if answer["category"][0] in ["yes", "no"]:  # category is list with one element
            _lowercase	       =[]
            _lowercase	       =[]
            _lowercase	       =input_ids[:q_len]
            _lowercase	       =range(A__      ,   len(A__     )      ,   max_length - doc_stride     )
            for i in doc_start_indices:
                  _lowercase	       =i + max_length - q_len
                  _lowercase	       =input_ids[i:end_index]
                  inputs.append(q_indices + slice     )
                  category.append(answer['category'][0]     )
                  if slice[-1] == tokenizer.sep_token_id:
                        break

            return {
                "example_id": example["id"],
                "input_ids": inputs,
                "labels": {
                    "start_token": [-100] * len(A__     ),
                    "end_token": [-100] * len(A__     ),
                    "category": category,
                },
            }

      _lowercase	       =out['context'].split()
      _lowercase	       =splitted_context[answer['end_token']]
      _lowercase	       =len(
          tokenizer(
              ' '.join(splitted_context[: answer['start_token']]     )      ,   add_special_tokens=A__      ,   ).input_ids     )
      _lowercase	       =len(
          tokenizer(' '.join(splitted_context[: answer['end_token']]     )      ,   add_special_tokens=A__     ).input_ids     )

      answer["start_token"] += q_len
      answer["end_token"] += q_len

      # fixing end token
      _lowercase	       =len(tokenizer(A__      ,   add_special_tokens=A__     ).input_ids     )
      if num_sub_tokens > 1:
            answer["end_token"] += num_sub_tokens - 1

      _lowercase	       =input_ids[answer['start_token'] : answer['end_token'] + 1]  # right & left are inclusive
      _lowercase	       =answer['start_token']
      _lowercase	       =answer['end_token']

      if assertion:
            _lowercase	       =tokenizer.decode(A__     )
            if answer["span"] != new:
                  print('ISSUE IN TOKENIZATION'     )
                  print('OLD:'      ,   answer['span']     )
                  print('NEW:'      ,   A__      ,   end='\n\n'     )

      if len(A__     ) <= max_length:
            return {
                "example_id": example["id"],
                "input_ids": [input_ids],
                "labels": {
                    "start_token": [answer["start_token"]],
                    "end_token": [answer["end_token"]],
                    "category": answer["category"],
                },
            }

      _lowercase	       =input_ids[:q_len]
      _lowercase	       =range(A__      ,   len(A__     )      ,   max_length - doc_stride     )

      _lowercase	       =[]
      _lowercase	       =[]
      _lowercase	       =[]
      _lowercase	       =[]  # null, yes, no, long, short
      for i in doc_start_indices:
            _lowercase	       =i + max_length - q_len
            _lowercase	       =input_ids[i:end_index]
            inputs.append(q_indices + slice     )
            assert len(inputs[-1]     ) <= max_length, "Issue in truncating length"

            if start_token >= i and end_token <= end_index - 1:
                  _lowercase	       =start_token - i + q_len
                  _lowercase	       =end_token - i + q_len
                  answers_category.append(answer['category'][0]     )  # ["short"] -> "short"
            else:
                  _lowercase	       =-100
                  _lowercase	       =-100
                  answers_category.append('null'     )
            _lowercase	       =inputs[-1][start_token : end_token + 1]

            answers_start_token.append(A__     )
            answers_end_token.append(A__     )
            if assertion:
                  if new != old and new != [tokenizer.cls_token_id]:
                        print('ISSUE in strided for ID:'      ,   example['id']     )
                        print('New:'      ,   tokenizer.decode(A__     )     )
                        print('Old:'      ,   tokenizer.decode(A__     )      ,   end='\n\n'     )
            if slice[-1] == tokenizer.sep_token_id:
                  break

      return {
          "example_id": example["id"],
          "input_ids": inputs,
          "labels": {
              "start_token": answers_start_token,
              "end_token": answers_end_token,
              "category": answers_category,
          },
      }
def      a      (			A__	:   str      ,   A__	:   Union[str, Any]      ,   A__	:   Any=2048      ,   A__	:   List[str]=4096      ,   A__	:   List[Any]=False     )    ->      Tuple:





      """simple docstring"""



      _lowercase	       =get_strided_contexts_and_ans(
          A__      ,   A__      ,   doc_stride=A__      ,   max_length=A__      ,   assertion=A__      ,   )

      return example
def      a      (			A__	:   Optional[int]      ,   A__	:   Any     )    ->      Any:





      """simple docstring"""



      with jsonlines.open(A__      ,   'a'     ) as writer:
            for example in tqdm(A__      ,   total=len(A__     )      ,   desc='Saving samples ... '     ):
                  _lowercase	       =example['labels']
                  for ids, start, end, cat in zip(
                      example['input_ids']      ,   labels['start_token']      ,   labels['end_token']      ,   labels['category']      ,   ):
                        if start == -1 and end == -1:
                              continue  # leave waste samples with no answer
                        if cat == "null" and np.random.rand() < 0.6:
                              continue  # removing 50 % samples
                        writer.write(
                            {
                                'input_ids': ids,
                                'start_token': start,
                                'end_token': end,
                                'category': CATEGORY_MAPPING[cat],
                            }     )


if __name__ == "__main__":
  from datasets import load_dataset

  from transformers import BigBirdTokenizer

  lowercase_									=   load_dataset('natural_questions')
  lowercase_									=   BigBirdTokenizer.from_pretrained('google/bigbird-roberta-base')

  lowercase_									=   data['train' if PROCESS_TRAIN == 'true' else 'validation']

  lowercase_									=   {
      'tokenizer': tokenizer,
      'doc_stride': DOC_STRIDE,
      'max_length': MAX_LENGTH,
      'assertion': False,
  }
  lowercase_									=   data.map(prepare_inputs, fn_kwargs=fn_kwargs)
  lowercase_									=   data.remove_columns(['annotations', 'document', 'id', 'question'])
  print(data)

  np.random.seed(SEED)
  lowercase_									=   'nq-training.jsonl' if PROCESS_TRAIN == 'true' else 'nq-validation.jsonl'
  save_to_disk(data, file_name=cache_file_name)





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


