





'''simple docstring'''




import unittest

from parameterized import parameterized

from transformers import LlamaConfig, is_torch_available, set_seed
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch

  from transformers import LlamaForCausalLM, LlamaForSequenceClassification, LlamaModel, LlamaTokenizer







class   A_	:




 def __init__(					self       :    Any	,			snake_case_       :    Tuple	,			snake_case_       :    Optional[Any]=1_3	,			snake_case_       :    Optional[int]=7	,			snake_case_       :    Optional[Any]=True	,			snake_case_       :    Dict=True	,			snake_case_       :    Dict=False	,			snake_case_       :    List[str]=True	,			snake_case_       :    Any=9_9	,			snake_case_       :    int=3_2	,			snake_case_       :    Any=5	,			snake_case_       :    Any=4	,			snake_case_       :    Tuple=3_7	,			snake_case_       :    Any="gelu"	,			snake_case_       :    List[str]=0.1	,			snake_case_       :    Any=0.1	,			snake_case_       :    int=5_1_2	,			snake_case_       :    Dict=1_6	,			snake_case_       :    int=2	,			snake_case_       :    str=0.0_2	,			snake_case_       :    str=3	,			snake_case_       :    Optional[Any]=4	,			snake_case_       :    str=None	,			):
      _UpperCAmelCase     							=   parent
      _UpperCAmelCase     							=   batch_size
      _UpperCAmelCase     							=   seq_length
      _UpperCAmelCase     							=   is_training
      _UpperCAmelCase     							=   use_input_mask
      _UpperCAmelCase     							=   use_token_type_ids
      _UpperCAmelCase     							=   use_labels
      _UpperCAmelCase     							=   vocab_size
      _UpperCAmelCase     							=   hidden_size
      _UpperCAmelCase     							=   num_hidden_layers
      _UpperCAmelCase     							=   num_attention_heads
      _UpperCAmelCase     							=   intermediate_size
      _UpperCAmelCase     							=   hidden_act
      _UpperCAmelCase     							=   hidden_dropout_prob
      _UpperCAmelCase     							=   attention_probs_dropout_prob
      _UpperCAmelCase     							=   max_position_embeddings
      _UpperCAmelCase     							=   type_vocab_size
      _UpperCAmelCase     							=   type_sequence_label_size
      _UpperCAmelCase     							=   initializer_range
      _UpperCAmelCase     							=   num_labels
      _UpperCAmelCase     							=   num_choices
      _UpperCAmelCase     							=   scope




 def      lowercase	(					self       :    str    ):
      _UpperCAmelCase     							=   ids_tensor([self.batch_size, self.seq_length]	,			self.vocab_size    )

      _UpperCAmelCase     							=   None
      if self.use_input_mask:
           _UpperCAmelCase     							=   random_attention_mask([self.batch_size, self.seq_length]    )

      _UpperCAmelCase     							=   None
      if self.use_token_type_ids:
           _UpperCAmelCase     							=   ids_tensor([self.batch_size, self.seq_length]	,			self.type_vocab_size    )

      _UpperCAmelCase     							=   None
      _UpperCAmelCase     							=   None
      _UpperCAmelCase     							=   None
      if self.use_labels:
           _UpperCAmelCase     							=   ids_tensor([self.batch_size]	,			self.type_sequence_label_size    )
           _UpperCAmelCase     							=   ids_tensor([self.batch_size, self.seq_length]	,			self.num_labels    )
           _UpperCAmelCase     							=   ids_tensor([self.batch_size]	,			self.num_choices    )

      _UpperCAmelCase     							=   self.get_config()

      return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels




 def      lowercase	(					self       :    Dict    ):
      return LlamaConfig(
          vocab_size=self.vocab_size	,			hidden_size=self.hidden_size	,			num_hidden_layers=self.num_hidden_layers	,			num_attention_heads=self.num_attention_heads	,			intermediate_size=self.intermediate_size	,			hidden_act=self.hidden_act	,			hidden_dropout_prob=self.hidden_dropout_prob	,			attention_probs_dropout_prob=self.attention_probs_dropout_prob	,			max_position_embeddings=self.max_position_embeddings	,			type_vocab_size=self.type_vocab_size	,			is_decoder=snake_case_	,			initializer_range=self.initializer_range	,			)




 def      lowercase	(					self       :    Any	,			snake_case_       :    Union[str, Any]	,			snake_case_       :    Union[str, Any]	,			snake_case_       :    Optional[int]	,			snake_case_       :    Dict	,			snake_case_       :    List[str]	,			snake_case_       :    List[Any]	,			snake_case_       :    Any    ):
      _UpperCAmelCase     							=   LlamaModel(config=snake_case_    )
      model.to(snake_case_    )
      model.eval()
      _UpperCAmelCase     							=   model(snake_case_	,			attention_mask=snake_case_    )
      _UpperCAmelCase     							=   model(snake_case_    )
      self.parent.assertEqual(result.last_hidden_state.shape	,			(self.batch_size, self.seq_length, self.hidden_size)    )




 def      lowercase	(					self       :    List[str]	,			snake_case_       :    Dict	,			snake_case_       :    Tuple	,			snake_case_       :    Union[str, Any]	,			snake_case_       :    Optional[Any]	,			snake_case_       :    str	,			snake_case_       :    Optional[Any]	,			snake_case_       :    Optional[int]	,			snake_case_       :    Dict	,			snake_case_       :    Optional[Any]	,			):
      _UpperCAmelCase     							=   True
      _UpperCAmelCase     							=   LlamaModel(snake_case_    )
      model.to(snake_case_    )
      model.eval()
      _UpperCAmelCase     							=   model(
          snake_case_	,			attention_mask=snake_case_	,			encoder_hidden_states=snake_case_	,			encoder_attention_mask=snake_case_	,			)
      _UpperCAmelCase     							=   model(
          snake_case_	,			attention_mask=snake_case_	,			encoder_hidden_states=snake_case_	,			)
      _UpperCAmelCase     							=   model(snake_case_	,			attention_mask=snake_case_    )
      self.parent.assertEqual(result.last_hidden_state.shape	,			(self.batch_size, self.seq_length, self.hidden_size)    )




 def      lowercase	(					self       :    Optional[int]	,			snake_case_       :    Union[str, Any]	,			snake_case_       :    Union[str, Any]	,			snake_case_       :    Union[str, Any]	,			snake_case_       :    Dict	,			snake_case_       :    Optional[int]	,			snake_case_       :    Optional[Any]	,			snake_case_       :    Tuple	,			snake_case_       :    str	,			snake_case_       :    Union[str, Any]	,			):
      _UpperCAmelCase     							=   LlamaForCausalLM(config=snake_case_    )
      model.to(snake_case_    )
      model.eval()
      _UpperCAmelCase     							=   model(snake_case_	,			attention_mask=snake_case_	,			labels=snake_case_    )
      self.parent.assertEqual(result.logits.shape	,			(self.batch_size, self.seq_length, self.vocab_size)    )




 def      lowercase	(					self       :    int	,			snake_case_       :    Dict	,			snake_case_       :    Union[str, Any]	,			snake_case_       :    List[Any]	,			snake_case_       :    Any	,			snake_case_       :    str	,			snake_case_       :    List[Any]	,			snake_case_       :    Union[str, Any]	,			snake_case_       :    Tuple	,			snake_case_       :    str	,			):
      _UpperCAmelCase     							=   True
      _UpperCAmelCase     							=   True
      _UpperCAmelCase     							=   LlamaForCausalLM(config=snake_case_    )
      model.to(snake_case_    )
      model.eval()

      # first forward pass
      _UpperCAmelCase     							=   model(
          snake_case_	,			attention_mask=snake_case_	,			encoder_hidden_states=snake_case_	,			encoder_attention_mask=snake_case_	,			use_cache=snake_case_	,			)
      _UpperCAmelCase     							=   outputs.past_key_values

      # create hypothetical multiple next token and extent to next_input_ids
      _UpperCAmelCase     							=   ids_tensor((self.batch_size, 3)	,			config.vocab_size    )
      _UpperCAmelCase     							=   ids_tensor((self.batch_size, 3)	,			vocab_size=2    )

      # append to next input_ids and
      _UpperCAmelCase     							=   torch.cat([input_ids, next_tokens]	,			dim=-1    )
      _UpperCAmelCase     							=   torch.cat([input_mask, next_mask]	,			dim=-1    )

      _UpperCAmelCase     							=   model(
          snake_case_	,			attention_mask=snake_case_	,			encoder_hidden_states=snake_case_	,			encoder_attention_mask=snake_case_	,			output_hidden_states=snake_case_	,			)["hidden_states"][0]
      _UpperCAmelCase     							=   model(
          snake_case_	,			attention_mask=snake_case_	,			encoder_hidden_states=snake_case_	,			encoder_attention_mask=snake_case_	,			past_key_values=snake_case_	,			output_hidden_states=snake_case_	,			)["hidden_states"][0]

      # select random slice
      _UpperCAmelCase     							=   ids_tensor((1,)	,			output_from_past.shape[-1]    ).item()
      _UpperCAmelCase     							=   output_from_no_past[:, -3:, random_slice_idx].detach()
      _UpperCAmelCase     							=   output_from_past[:, :, random_slice_idx].detach()

      self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]    )

      # test that outputs are equal for slice
      self.parent.assertTrue(torch.allclose(snake_case_	,			snake_case_	,			atol=1e-3    )    )




 def      lowercase	(					self       :    Optional[int]    ):
      _UpperCAmelCase     							=   self.prepare_config_and_inputs()
      (
          (
          _UpperCAmelCase
      )							,  (
          _UpperCAmelCase
      )							,  (
          _UpperCAmelCase
      )							,  (
          _UpperCAmelCase
      )							,  (
          _UpperCAmelCase
      )							,  (
          _UpperCAmelCase
      )							,  (
          _UpperCAmelCase
      )							,  
      )     							=   config_and_inputs
      _UpperCAmelCase     							=   {"input_ids": input_ids, "attention_mask": input_mask}
      return config, inputs_dict







@require_torch
class   A_	(  lowerCAmelCase_							,						lowerCAmelCase_							,						lowerCAmelCase_							,						unittest.TestCase			):
 _lowerCamelCase  : Dict           =							(LlamaModel, LlamaForCausalLM, LlamaForSequenceClassification) if is_torch_available() else ()
 _lowerCamelCase  : List[Any]           =							(LlamaForCausalLM,) if is_torch_available() else ()
 _lowerCamelCase  : List[str]           =							(
     {
         """feature-extraction""": LlamaModel,
         """text-classification""": LlamaForSequenceClassification,
         """text-generation""": LlamaForCausalLM,
         """zero-shot""": LlamaForSequenceClassification,
     }
     if is_torch_available()
     else {}
 )
 _lowerCamelCase  : Dict           =							False
 _lowerCamelCase  : List[str]           =							False




 def      lowercase	(					self       :    List[Any]    ):
      _UpperCAmelCase     							=   LlamaModelTester(self    )
      _UpperCAmelCase     							=   ConfigTester(self	,			config_class=snake_case_	,			hidden_size=3_7    )




 def      lowercase	(					self       :    Optional[Any]    ):
      self.config_tester.run_common_tests()




 def      lowercase	(					self       :    Optional[Any]    ):
      _UpperCAmelCase     							=   self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*snake_case_    )




 def      lowercase	(					self       :    int    ):
      _UpperCAmelCase     							=   self.model_tester.prepare_config_and_inputs()
      for type in ["absolute", "relative_key", "relative_key_query"]:
           _UpperCAmelCase     							=   type
           self.model_tester.create_and_check_model(*snake_case_    )




 def      lowercase	(					self       :    Optional[int]    ):
      _UpperCAmelCase							,  _UpperCAmelCase     							=   self.model_tester.prepare_config_and_inputs_for_common()
      _UpperCAmelCase     							=   3
      _UpperCAmelCase     							=   input_dict["input_ids"]
      _UpperCAmelCase     							=   input_ids.ne(1    ).to(snake_case_    )
      _UpperCAmelCase     							=   ids_tensor([self.model_tester.batch_size]	,			self.model_tester.type_sequence_label_size    )
      _UpperCAmelCase     							=   LlamaForSequenceClassification(snake_case_    )
      model.to(snake_case_    )
      model.eval()
      _UpperCAmelCase     							=   model(snake_case_	,			attention_mask=snake_case_	,			labels=snake_case_    )
      self.assertEqual(result.logits.shape	,			(self.model_tester.batch_size, self.model_tester.num_labels)    )




 def      lowercase	(					self       :    Any    ):
      _UpperCAmelCase							,  _UpperCAmelCase     							=   self.model_tester.prepare_config_and_inputs_for_common()
      _UpperCAmelCase     							=   3
      _UpperCAmelCase     							=   "single_label_classification"
      _UpperCAmelCase     							=   input_dict["input_ids"]
      _UpperCAmelCase     							=   input_ids.ne(1    ).to(snake_case_    )
      _UpperCAmelCase     							=   ids_tensor([self.model_tester.batch_size]	,			self.model_tester.type_sequence_label_size    )
      _UpperCAmelCase     							=   LlamaForSequenceClassification(snake_case_    )
      model.to(snake_case_    )
      model.eval()
      _UpperCAmelCase     							=   model(snake_case_	,			attention_mask=snake_case_	,			labels=snake_case_    )
      self.assertEqual(result.logits.shape	,			(self.model_tester.batch_size, self.model_tester.num_labels)    )




 def      lowercase	(					self       :    List[str]    ):
      _UpperCAmelCase							,  _UpperCAmelCase     							=   self.model_tester.prepare_config_and_inputs_for_common()
      _UpperCAmelCase     							=   3
      _UpperCAmelCase     							=   "multi_label_classification"
      _UpperCAmelCase     							=   input_dict["input_ids"]
      _UpperCAmelCase     							=   input_ids.ne(1    ).to(snake_case_    )
      _UpperCAmelCase     							=   ids_tensor(
          [self.model_tester.batch_size, config.num_labels]	,			self.model_tester.type_sequence_label_size    ).to(torch.float    )
      _UpperCAmelCase     							=   LlamaForSequenceClassification(snake_case_    )
      model.to(snake_case_    )
      model.eval()
      _UpperCAmelCase     							=   model(snake_case_	,			attention_mask=snake_case_	,			labels=snake_case_    )
      self.assertEqual(result.logits.shape	,			(self.model_tester.batch_size, self.model_tester.num_labels)    )




 @unittest.skip("LLaMA buffers include complex numbers, which breaks this test"    )
 def      lowercase	(					self       :    Tuple    ):
      pass




 @parameterized.expand([("linear",), ("dynamic",)]    )
 def      lowercase	(					self       :    Optional[int]	,			snake_case_       :    Optional[int]    ):
      _UpperCAmelCase							,  _UpperCAmelCase     							=   self.model_tester.prepare_config_and_inputs_for_common()
      _UpperCAmelCase     							=   ids_tensor([1, 1_0]	,			config.vocab_size    )
      _UpperCAmelCase     							=   ids_tensor([1, int(config.max_position_embeddings * 1.5    )]	,			config.vocab_size    )

      set_seed(4_2    )  # Fixed seed at init time so the two models get the same random weights
      _UpperCAmelCase     							=   LlamaModel(snake_case_    )
      original_model.to(snake_case_    )
      original_model.eval()
      _UpperCAmelCase     							=   original_model(snake_case_    ).last_hidden_state
      _UpperCAmelCase     							=   original_model(snake_case_    ).last_hidden_state

      set_seed(4_2    )  # Fixed seed at init time so the two models get the same random weights
      _UpperCAmelCase     							=   {"type": scaling_type, "factor": 1_0.0}
      _UpperCAmelCase     							=   LlamaModel(snake_case_    )
      scaled_model.to(snake_case_    )
      scaled_model.eval()
      _UpperCAmelCase     							=   scaled_model(snake_case_    ).last_hidden_state
      _UpperCAmelCase     							=   scaled_model(snake_case_    ).last_hidden_state

      # Dynamic scaling does not change the RoPE embeddings until it receives an input longer than the original
      # maximum sequence length, so the outputs for the short input should match.
      if scaling_type == "dynamic":
           self.assertTrue(torch.allclose(snake_case_	,			snake_case_	,			atol=1e-5    )    )
      else:
           self.assertFalse(torch.allclose(snake_case_	,			snake_case_	,			atol=1e-5    )    )

      # The output should be different for long inputs
      self.assertFalse(torch.allclose(snake_case_	,			snake_case_	,			atol=1e-5    )    )







@require_torch
class   A_	(  unittest.TestCase			):




 @unittest.skip("Logits are not exactly the same, once we fix the instabalities somehow, will update!"    )
 @slow
 def      lowercase	(					self       :    Any    ):
      _UpperCAmelCase     							=   [1, 3_0_6, 4_6_5_8, 2_7_8, 6_5_9_3, 3_1_0, 2_8_3_4, 3_3_8]
      _UpperCAmelCase     							=   LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf"	,			device_map="auto"    )
      _UpperCAmelCase     							=   model(torch.tensor([input_ids]    )    )
      # Expected mean on dim = -1
      _UpperCAmelCase     							=   torch.tensor([[-6.6_5_5_0, -4.1_2_2_7, -4.9_8_5_9, -3.2_4_0_6, 0.8_2_6_2, -3.0_0_3_3, 1.2_9_6_4, -3.3_6_9_9]]    )
      torch.testing.assert_close(out.mean(-1    )	,			snake_case_	,			atol=1e-2	,			rtol=1e-2    )
      # slicing logits[0, 0, 0:30]
      # fmt: off
      _UpperCAmelCase     							=   torch.tensor([-1_2.8_2_8_1, -7.4_4_5_3, -0.4_6_3_9, -8.0_6_2_5, -7.2_5_0_0, -8.0_0_0_0, -6.4_8_8_3, -7.7_6_9_5, -7.8_4_3_8, -7.0_3_1_2, -6.2_1_8_8, -7.1_3_2_8, -1.8_4_9_6, 1.9_9_6_1, -8.6_2_5_0, -6.7_2_2_7, -1_2.8_2_8_1, -6.9_4_9_2, -7.0_7_4_2, -7.7_8_5_2, -7.5_8_2_0, -7.9_0_6_2, -6.9_3_7_5, -7.9_8_0_5, -8.3_4_3_8, -8.1_5_6_2, -8.0_4_6_9, -7.6_2_5_0, -7.7_4_2_2, -7.3_3_9_8,]    )
      # fmt: on
      torch.testing.assert_close(out[0, 0, :3_0]	,			snake_case_	,			atol=1e-5	,			rtol=1e-5    )




 @unittest.skip("Logits are not exactly the same, once we fix the instabalities somehow, will update!"    )
 @slow
 def      lowercase	(					self       :    Union[str, Any]    ):
      _UpperCAmelCase     							=   [1, 3_0_6, 4_6_5_8, 2_7_8, 6_5_9_3, 3_1_0, 2_8_3_4, 3_3_8]
      _UpperCAmelCase     							=   LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-13b-hf"	,			device_map="auto"    )
      _UpperCAmelCase     							=   model(torch.tensor(snake_case_    )    )
      # Expected mean on dim = -1
      _UpperCAmelCase     							=   torch.tensor([[-2.0_6_2_2, -1.2_7_9_4, -1.1_6_3_8, -0.9_7_8_8, -1.4_6_0_3, -1.0_2_3_8, -1.7_8_9_3, -1.4_4_1_1]]    )
      torch.testing.assert_close(out.mean(-1    )	,			snake_case_	,			atol=1e-2	,			rtol=1e-2    )
      # slicing logits[0, 0, 0:30]
      # fmt: off
      _UpperCAmelCase     							=   torch.tensor([-8.1_4_0_6, -8.0_5_4_7, 2.7_4_6_1, -1.2_3_4_4, -0.1_4_4_8, -1.8_2_6_2, -1.0_0_2_0, -1.8_1_5_4, -1.6_8_9_5, -1.8_5_1_6, -2.3_5_7_4, -0.9_2_7_7, 3.7_5_9_8, 6.5_7_4_2, -1.2_9_9_8, -0.1_1_7_7, -8.1_4_0_6, -2.9_6_8_8, -2.9_1_9_9, -3.1_6_9_9, -3.5_2_5_4, -2.3_5_5_5, -2.7_9_8_8, -3.4_1_4_1, -2.8_2_6_2, -4.5_1_9_5, -3.3_3_7_9, -3.3_1_6_4, -2.7_8_3_2, -3.0_2_7_3]    )
      # fmt: on
      torch.testing.assert_close(out[0, 0, :3_0]	,			snake_case_	,			atol=1e-5	,			rtol=1e-5    )




 @unittest.skip("Logits are not exactly the same, once we fix the instabalities somehow, will update!"    )
 @slow
 def      lowercase	(					self       :    Any    ):
      _UpperCAmelCase     							=   [1, 3_0_6, 4_6_5_8, 2_7_8, 6_5_9_3, 3_1_0, 2_8_3_4, 3_3_8]
      _UpperCAmelCase     							=   LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-13b-chat-hf"	,			device_map="auto"    )
      _UpperCAmelCase     							=   model(torch.tensor(snake_case_    )    )
      # Expected mean on dim = -1
      _UpperCAmelCase     							=   torch.tensor([[-0.8_5_6_2, -1.8_5_2_0, -0.7_5_5_1, -0.4_1_6_2, -1.5_1_6_1, -1.2_0_3_8, -2.4_8_2_3, -2.3_2_5_4]]    )
      torch.testing.assert_close(out.mean(-1    )	,			snake_case_	,			atol=1e-2	,			rtol=1e-2    )
      # slicing logits[0, 0, 0:30]
      # fmt: off
      _UpperCAmelCase     							=   torch.tensor([-2.2_2_2_7, 4.8_8_2_8, 0.9_0_2_3, -0.4_5_7_8, -0.7_8_7_1, -0.1_0_3_3, -0.6_2_2_1, -0.5_7_8_6, -0.7_8_0_3, -1.0_6_7_4, -1.2_9_2_0, -0.1_5_7_0, 0.8_0_0_8, 2.0_7_2_3, -0.9_4_9_7, 0.2_7_7_1, -2.2_2_2_7, -0.7_6_1_2, -1.4_3_4_6, -1.2_0_6_1, -1.6_4_2_6, -0.3_0_0_0, -0.7_1_3_9, -1.1_9_3_4, -1.8_6_9_1, -1.6_9_7_3, -1.5_9_4_7, -1.2_7_0_5, -0.3_5_2_3, -0.5_5_1_3]    )
      # fmt: on
      torch.testing.assert_close(out.mean(-1    )	,			snake_case_	,			atol=1e-2	,			rtol=1e-2    )




 @unittest.skip(
     "Logits are not exactly the same, once we fix the instabalities somehow, will update! Also it is gonna be a `too_slow` test"    )
 @slow
 def      lowercase	(					self       :    List[Any]    ):
      _UpperCAmelCase     							=   [1, 3_0_6, 4_6_5_8, 2_7_8, 6_5_9_3, 3_1_0, 2_8_3_4, 3_3_8]
      _UpperCAmelCase     							=   LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-70b-hf"	,			device_map="auto"    )
      _UpperCAmelCase     							=   model(torch.tensor(snake_case_    )    )

      _UpperCAmelCase     							=   torch.tensor(
          [[-4.2_3_2_7, -3.3_3_6_0, -4.6_6_6_5, -4.7_6_3_1, -1.8_1_8_0, -3.4_1_7_0, -1.4_2_1_1, -3.1_8_1_0]]	,			dtype=torch.floataa    )
      torch.testing.assert_close(out.mean(-1    )	,			snake_case_	,			atol=1e-2	,			rtol=1e-2    )
      # fmt: off
      _UpperCAmelCase     							=   torch.tensor([-9.4_9_2_2, -3.9_5_5_1, 1.7_9_9_8, -5.6_7_5_8, -5.1_0_5_5, -5.8_9_8_4, -4.8_3_2_0, -6.8_0_8_6, -6.5_3_9_1, -5.6_1_7_2, -5.5_8_2_0, -5.5_3_5_2, 1.7_8_8_1, 3.6_2_8_9, -6.5_1_1_7, -3.4_7_8_5, -9.5_0_0_0, -6.0_3_5_2, -6.8_1_2_5, -6.0_1_9_5, -6.6_8_3_6, -5.4_7_2_7, -6.2_8_1_2, -6.0_3_9_1, -7.3_3_9_8, -7.4_2_9_7, -7.4_8_4_4, -6.5_8_2_0, -5.8_7_8_9, -5.5_3_1_2]    )
      # fmt: on
      torch.testing.assert_close(out[0, 0, :3_0]	,			snake_case_	,			atol=1e-5	,			rtol=1e-5    )




 @unittest.skip("Model is curently gated"    )
 @slow
 def      lowercase	(					self       :    Optional[Any]    ):
      _UpperCAmelCase     							=   "Simply put, the theory of relativity states that 1) the laws of physics are the same everywhere in the universe and 2) the passage of time and the length of objects can vary depending on the observer\'s frame of reference.\n\nThe first part of the theory, that the laws of physics are the same everywhere, is known as the \"princi"
      _UpperCAmelCase     							=   "Simply put, the theory of relativity states that "
      _UpperCAmelCase     							=   LlamaTokenizer.from_pretrained("meta-llama/Llama-2-13b-chat-hf"    )
      _UpperCAmelCase     							=   tokenizer.encode(snake_case_	,			return_tensors="pt"    )
      _UpperCAmelCase     							=   LlamaForCausalLM.from_pretrained(
          "meta-llama/Llama-2-13b-chat-hf"	,			device_map="sequential"	,			use_safetensors=snake_case_    )

      # greedy generation outputs
      _UpperCAmelCase     							=   model.generate(snake_case_	,			max_new_tokens=6_4	,			top_p=snake_case_	,			temperature=1	,			do_sample=snake_case_    )
      _UpperCAmelCase     							=   tokenizer.decode(generated_ids[0]	,			skip_special_tokens=snake_case_    )
      self.assertEqual(snake_case_	,			snake_case_    )











'''simple docstring'''




import inspect
import tempfile
import unittest

from huggingface_hub import hf_hub_download

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

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


__SCREAMING_SNAKE_CASE     :Dict       =				1e-4

if is_torch_available():
  import torch

  from transformers import AutoformerConfig, AutoformerForPrediction, AutoformerModel
  from transformers.models.autoformer.modeling_autoformer import AutoformerDecoder, AutoformerEncoder







@require_torch
class   A_	:




 def __init__(					self       :    List[Any]	,			snake_case_       :    int	,			snake_case_       :    Dict=1_6	,			snake_case_       :    Dict=1_3	,			snake_case_       :    int=7	,			snake_case_       :    Any=1_4	,			snake_case_       :    int=1_0	,			snake_case_       :    Any=1_9	,			snake_case_       :    int=5	,			snake_case_       :    Any=4	,			snake_case_       :    Tuple=True	,			snake_case_       :    Optional[int]=1_6	,			snake_case_       :    List[str]=2	,			snake_case_       :    Any=4	,			snake_case_       :    List[Any]=4	,			snake_case_       :    Optional[Any]="gelu"	,			snake_case_       :    Optional[int]=0.1	,			snake_case_       :    Union[str, Any]=0.1	,			snake_case_       :    Tuple=[1, 2, 3, 4, 5]	,			snake_case_       :    str=2_5	,			snake_case_       :    Any=5	,			):
      _UpperCAmelCase     							=   d_model
      _UpperCAmelCase     							=   parent
      _UpperCAmelCase     							=   batch_size
      _UpperCAmelCase     							=   prediction_length
      _UpperCAmelCase     							=   context_length
      _UpperCAmelCase     							=   cardinality
      _UpperCAmelCase     							=   num_time_features
      _UpperCAmelCase     							=   lags_sequence
      _UpperCAmelCase     							=   embedding_dimension
      _UpperCAmelCase     							=   is_training
      _UpperCAmelCase     							=   hidden_size
      _UpperCAmelCase     							=   num_hidden_layers
      _UpperCAmelCase     							=   num_attention_heads
      _UpperCAmelCase     							=   intermediate_size
      _UpperCAmelCase     							=   hidden_act
      _UpperCAmelCase     							=   hidden_dropout_prob
      _UpperCAmelCase     							=   attention_probs_dropout_prob

      _UpperCAmelCase     							=   context_length
      _UpperCAmelCase     							=   prediction_length + label_length
      _UpperCAmelCase     							=   label_length

      _UpperCAmelCase     							=   moving_average
      _UpperCAmelCase     							=   autocorrelation_factor




 def      lowercase	(					self       :    Union[str, Any]    ):
      return AutoformerConfig(
          d_model=self.d_model	,			encoder_layers=self.num_hidden_layers	,			decoder_layers=self.num_hidden_layers	,			encoder_attention_heads=self.num_attention_heads	,			decoder_attention_heads=self.num_attention_heads	,			encoder_ffn_dim=self.intermediate_size	,			decoder_ffn_dim=self.intermediate_size	,			dropout=self.hidden_dropout_prob	,			attention_dropout=self.attention_probs_dropout_prob	,			prediction_length=self.prediction_length	,			context_length=self.context_length	,			label_length=self.label_length	,			lags_sequence=self.lags_sequence	,			num_time_features=self.num_time_features	,			num_static_categorical_features=1	,			cardinality=[self.cardinality]	,			embedding_dimension=[self.embedding_dimension]	,			moving_average=self.moving_average	,			)




 def      lowercase	(					self       :    int	,			snake_case_       :    Optional[Any]    ):
      _UpperCAmelCase     							=   config.context_length + max(config.lags_sequence    )

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

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

      _UpperCAmelCase     							=   {
          "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      lowercase	(					self       :    List[Any]    ):
      _UpperCAmelCase     							=   self.get_config()
      _UpperCAmelCase     							=   self.prepare_autoformer_inputs_dict(snake_case_    )
      return config, inputs_dict




 def      lowercase	(					self       :    List[str]    ):
      _UpperCAmelCase							,  _UpperCAmelCase     							=   self.prepare_config_and_inputs()
      return config, inputs_dict




 def      lowercase	(					self       :    Optional[Any]	,			snake_case_       :    int	,			snake_case_       :    Optional[int]    ):
      _UpperCAmelCase     							=   AutoformerModel(config=snake_case_    ).to(snake_case_    ).eval()
      _UpperCAmelCase     							=   model(**snake_case_    )

      _UpperCAmelCase     							=   outputs.encoder_last_hidden_state
      _UpperCAmelCase     							=   outputs.last_hidden_state

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

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

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

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

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

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

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







@require_torch
class   A_	(  lowerCAmelCase_							,						lowerCAmelCase_							,						unittest.TestCase			):
 _lowerCamelCase  : List[Any]           =							(AutoformerModel, AutoformerForPrediction) if is_torch_available() else ()
 _lowerCamelCase  : Tuple           =							(AutoformerForPrediction,) if is_torch_available() else ()
 _lowerCamelCase  : List[Any]           =							{"""feature-extraction""": AutoformerModel} if is_torch_available() else {}
 _lowerCamelCase  : Optional[Any]           =							False
 _lowerCamelCase  : Tuple           =							False
 _lowerCamelCase  : int           =							False
 _lowerCamelCase  : Optional[Any]           =							False
 _lowerCamelCase  : Optional[Any]           =							False
 _lowerCamelCase  : List[Any]           =							False




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




 def      lowercase	(					self       :    Optional[Any]    ):
      self.config_tester.run_common_tests()




 def      lowercase	(					self       :    Union[str, Any]    ):
      _UpperCAmelCase							,  _UpperCAmelCase     							=   self.model_tester.prepare_config_and_inputs()
      for model_class in self.all_model_classes:
           _UpperCAmelCase     							=   model_class(snake_case_    )

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




 def      lowercase	(					self       :    Optional[int]    ):
      _UpperCAmelCase     							=   self.model_tester.prepare_config_and_inputs_for_common()
      self.model_tester.check_encoder_decoder_model_standalone(*snake_case_    )




 @unittest.skip(reason="Model has no tokens embeddings"    )
 def      lowercase	(					self       :    Optional[int]    ):
      pass




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




 def      lowercase	(					self       :    List[str]    ):
      _UpperCAmelCase							,  _UpperCAmelCase     							=   self.model_tester.prepare_config_and_inputs_for_common()

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

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




 def      lowercase	(					self       :    Optional[int]    ):
      _UpperCAmelCase							,  _UpperCAmelCase     							=   self.model_tester.prepare_config_and_inputs_for_common()
      _UpperCAmelCase     							=   True

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

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

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

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

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

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

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

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

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





def 	UpperCAmelCase_							(					__lowercase							:      str="train-batch.pt"       )		->       List[str]:


     '''simple docstring'''



     _UpperCAmelCase     							=   hf_hub_download(repo_id="hf-internal-testing/tourism-monthly-batch"			,     filename=__lowercase			,     repo_type="dataset"       )
     _UpperCAmelCase     							=   torch.load(__lowercase			,     map_location=__lowercase       )
     return batch







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




 def      lowercase	(					self       :    Optional[int]    ):
      _UpperCAmelCase     							=   AutoformerModel.from_pretrained("huggingface/autoformer-tourism-monthly"    ).to(snake_case_    )
      _UpperCAmelCase     							=   prepare_batch()

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

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

      _UpperCAmelCase     							=   torch.tensor(
          [[0.3_5_9_3, -1.3_3_9_8, 0.6_3_3_0], [0.2_2_7_9, 1.5_3_9_6, -0.1_7_9_2], [0.0_4_5_0, 1.3_2_2_5, -0.2_3_3_5]]	,			device=snake_case_    )
      self.assertTrue(torch.allclose(output[0, :3, :3]	,			snake_case_	,			atol=snake_case_    )    )




 def      lowercase	(					self       :    Optional[Any]    ):
      _UpperCAmelCase     							=   AutoformerForPrediction.from_pretrained("huggingface/autoformer-tourism-monthly"    ).to(snake_case_    )
      _UpperCAmelCase     							=   prepare_batch("val-batch.pt"    )
      with torch.no_grad():
           _UpperCAmelCase     							=   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
      _UpperCAmelCase     							=   torch.Size((6_4, model.config.context_length, model.config.d_model)    )
      self.assertEqual(output.shape	,			snake_case_    )

      _UpperCAmelCase     							=   torch.tensor(
          [[-0.0_7_3_4, -0.9_0_3_6, 0.8_3_5_8], [4.7_1_8_6, 2.4_1_1_3, 1.9_5_8_1], [1.7_9_5_3, 2.3_5_5_8, 1.2_9_7_0]]	,			device=snake_case_    )
      self.assertTrue(torch.allclose(output[0, :3, :3]	,			snake_case_	,			atol=snake_case_    )    )




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

      _UpperCAmelCase     							=   torch.tensor([3_1_3_0.6_7_6_3, 4_0_5_6.5_2_9_3, 7_0_5_3.0_7_8_6]	,			device=snake_case_    )
      _UpperCAmelCase     							=   outputs.sequences.mean(dim=1    )
      self.assertTrue(torch.allclose(mean_prediction[0, -3:]	,			snake_case_	,			rtol=1e-1    )    )




