


import unittest

from parameterized import parameterized

from transformers import OpenLlamaConfig, is_torch_available, set_seed
from transformers.testing_utils import require_torch, 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 OpenLlamaForCausalLM, OpenLlamaForSequenceClassification, OpenLlamaModel



class 							_UpperCamelCase :

   '''simple docstring'''







   def __init__(  self  : Tuple       ,  _lowerCamelCase  : Optional[Any]       ,  _lowerCamelCase  : Any=1_3       ,  _lowerCamelCase  : List[str]=7       ,  _lowerCamelCase  : int=True       ,  _lowerCamelCase  : Union[str, Any]=True       ,  _lowerCamelCase  : Dict=False       ,  _lowerCamelCase  : List[Any]=True       ,  _lowerCamelCase  : Optional[Any]=9_9       ,  _lowerCamelCase  : List[Any]=3_2       ,  _lowerCamelCase  : Union[str, Any]=5       ,  _lowerCamelCase  : Union[str, Any]=4       ,  _lowerCamelCase  : str=3_7       ,  _lowerCamelCase  : str="gelu"       ,  _lowerCamelCase  : Optional[Any]=0.1       ,  _lowerCamelCase  : int=0.1       ,  _lowerCamelCase  : str=5_1_2       ,  _lowerCamelCase  : Tuple=1_6       ,  _lowerCamelCase  : int=2       ,  _lowerCamelCase  : Optional[Any]=0.02       ,  _lowerCamelCase  : Optional[int]=3       ,  _lowerCamelCase  : Union[str, Any]=4       ,  _lowerCamelCase  : List[str]=None       ,  ):






         '''simple docstring'''



         __lowerCamelCase     :      int       						=						parent
         __lowerCamelCase     :      Optional[Any]       						=						batch_size
         __lowerCamelCase     :      str       						=						seq_length
         __lowerCamelCase     :      List[str]       						=						is_training
         __lowerCamelCase     :      Dict       						=						use_input_mask
         __lowerCamelCase     :      Optional[Any]       						=						use_token_type_ids
         __lowerCamelCase     :      int       						=						use_labels
         __lowerCamelCase     :      int       						=						vocab_size
         __lowerCamelCase     :      Union[str, Any]       						=						hidden_size
         __lowerCamelCase     :      int       						=						num_hidden_layers
         __lowerCamelCase     :      Any       						=						num_attention_heads
         __lowerCamelCase     :      Any       						=						intermediate_size
         __lowerCamelCase     :      str       						=						hidden_act
         __lowerCamelCase     :      str       						=						hidden_dropout_prob
         __lowerCamelCase     :      List[Any]       						=						attention_probs_dropout_prob
         __lowerCamelCase     :      Dict       						=						max_position_embeddings
         __lowerCamelCase     :      int       						=						type_vocab_size
         __lowerCamelCase     :      Optional[int]       						=						type_sequence_label_size
         __lowerCamelCase     :      Dict       						=						initializer_range
         __lowerCamelCase     :      List[str]       						=						num_labels
         __lowerCamelCase     :      List[str]       						=						num_choices
         __lowerCamelCase     :      Optional[int]       						=						scope

   def   _snake_case    (  self  : Tuple							):






         '''simple docstring'''



         __lowerCamelCase     :      int       						=						ids_tensor([self.batch_size, self.seq_length]       ,  self.vocab_size							)

         __lowerCamelCase     :      Optional[int]       						=						None
         if self.use_input_mask:
               __lowerCamelCase     :      Dict       						=						random_attention_mask([self.batch_size, self.seq_length]							)

         __lowerCamelCase     :      int       						=						None
         if self.use_token_type_ids:
               __lowerCamelCase     :      List[str]       						=						ids_tensor([self.batch_size, self.seq_length]       ,  self.type_vocab_size							)

         __lowerCamelCase     :      Dict       						=						None
         __lowerCamelCase     :      Dict       						=						None
         __lowerCamelCase     :      Dict       						=						None
         if self.use_labels:
               __lowerCamelCase     :      List[str]       						=						ids_tensor([self.batch_size]       ,  self.type_sequence_label_size							)
               __lowerCamelCase     :      Union[str, Any]       						=						ids_tensor([self.batch_size, self.seq_length]       ,  self.num_labels							)
               __lowerCamelCase     :      Dict       						=						ids_tensor([self.batch_size]       ,  self.num_choices							)

         __lowerCamelCase     :      List[Any]       						=						self.get_config()

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

   def   _snake_case    (  self  : Union[str, Any]							):






         '''simple docstring'''



         return OpenLlamaConfig(
             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=_lowerCamelCase       ,  initializer_range=self.initializer_range       ,  use_stable_embedding=_lowerCamelCase       ,  )

   def   _snake_case    (  self  : str       ,  _lowerCamelCase  : Union[str, Any]       ,  _lowerCamelCase  : Union[str, Any]       ,  _lowerCamelCase  : Union[str, Any]       ,  _lowerCamelCase  : List[str]       ,  _lowerCamelCase  : Union[str, Any]       ,  _lowerCamelCase  : Union[str, Any]       ,  _lowerCamelCase  : str							):






         '''simple docstring'''



         __lowerCamelCase     :      Union[str, Any]       						=						OpenLlamaModel(config=_lowerCamelCase							)
         model.to(_lowerCamelCase							)
         model.eval()
         __lowerCamelCase     :      Dict       						=						model(_lowerCamelCase       ,  attention_mask=_lowerCamelCase							)
         __lowerCamelCase     :      Dict       						=						model(_lowerCamelCase							)
         self.parent.assertEqual(result.last_hidden_state.shape       ,  (self.batch_size, self.seq_length, self.hidden_size)							)

   def   _snake_case    (  self  : Optional[Any]       ,  _lowerCamelCase  : Optional[Any]       ,  _lowerCamelCase  : int       ,  _lowerCamelCase  : Tuple       ,  _lowerCamelCase  : Optional[Any]       ,  _lowerCamelCase  : Dict       ,  _lowerCamelCase  : List[str]       ,  _lowerCamelCase  : Union[str, Any]       ,  _lowerCamelCase  : Union[str, Any]       ,  _lowerCamelCase  : Optional[Any]       ,  ):






         '''simple docstring'''



         __lowerCamelCase     :      int       						=						True
         __lowerCamelCase     :      Any       						=						OpenLlamaModel(_lowerCamelCase							)
         model.to(_lowerCamelCase							)
         model.eval()
         __lowerCamelCase     :      List[str]       						=						model(
             _lowerCamelCase       ,  attention_mask=_lowerCamelCase       ,  encoder_hidden_states=_lowerCamelCase       ,  encoder_attention_mask=_lowerCamelCase       ,  )
         __lowerCamelCase     :      Any       						=						model(
             _lowerCamelCase       ,  attention_mask=_lowerCamelCase       ,  encoder_hidden_states=_lowerCamelCase       ,  )
         __lowerCamelCase     :      List[str]       						=						model(_lowerCamelCase       ,  attention_mask=_lowerCamelCase							)
         self.parent.assertEqual(result.last_hidden_state.shape       ,  (self.batch_size, self.seq_length, self.hidden_size)							)

   def   _snake_case    (  self  : List[Any]       ,  _lowerCamelCase  : List[str]       ,  _lowerCamelCase  : Optional[Any]       ,  _lowerCamelCase  : Any       ,  _lowerCamelCase  : List[Any]       ,  _lowerCamelCase  : List[str]       ,  _lowerCamelCase  : Any       ,  _lowerCamelCase  : Optional[Any]       ,  _lowerCamelCase  : Any       ,  _lowerCamelCase  : Tuple       ,  ):






         '''simple docstring'''



         __lowerCamelCase     :      Optional[Any]       						=						OpenLlamaForCausalLM(config=_lowerCamelCase							)
         model.to(_lowerCamelCase							)
         model.eval()
         __lowerCamelCase     :      Tuple       						=						model(_lowerCamelCase       ,  attention_mask=_lowerCamelCase       ,  labels=_lowerCamelCase							)
         self.parent.assertEqual(result.logits.shape       ,  (self.batch_size, self.seq_length, self.vocab_size)							)

   def   _snake_case    (  self  : List[Any]       ,  _lowerCamelCase  : Union[str, Any]       ,  _lowerCamelCase  : Optional[Any]       ,  _lowerCamelCase  : Tuple       ,  _lowerCamelCase  : List[str]       ,  _lowerCamelCase  : Dict       ,  _lowerCamelCase  : Optional[Any]       ,  _lowerCamelCase  : Dict       ,  _lowerCamelCase  : Tuple       ,  _lowerCamelCase  : Optional[Any]       ,  ):






         '''simple docstring'''



         __lowerCamelCase     :      Dict       						=						True
         __lowerCamelCase     :      List[str]       						=						True
         __lowerCamelCase     :      List[Any]       						=						OpenLlamaForCausalLM(config=_lowerCamelCase							)
         model.to(_lowerCamelCase							)
         model.eval()

         # first forward pass
         __lowerCamelCase     :      List[str]       						=						model(
             _lowerCamelCase       ,  attention_mask=_lowerCamelCase       ,  encoder_hidden_states=_lowerCamelCase       ,  encoder_attention_mask=_lowerCamelCase       ,  use_cache=_lowerCamelCase       ,  )
         __lowerCamelCase     :      int       						=						outputs.past_key_values

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

         # append to next input_ids and
         __lowerCamelCase     :      Tuple       						=						torch.cat([input_ids, next_tokens]       ,  dim=-1							)
         __lowerCamelCase     :      Optional[int]       						=						torch.cat([input_mask, next_mask]       ,  dim=-1							)

         __lowerCamelCase     :      Optional[int]       						=						model(
             _lowerCamelCase       ,  attention_mask=_lowerCamelCase       ,  encoder_hidden_states=_lowerCamelCase       ,  encoder_attention_mask=_lowerCamelCase       ,  output_hidden_states=_lowerCamelCase       ,  )["""hidden_states"""][0]
         __lowerCamelCase     :      Optional[int]       						=						model(
             _lowerCamelCase       ,  attention_mask=_lowerCamelCase       ,  encoder_hidden_states=_lowerCamelCase       ,  encoder_attention_mask=_lowerCamelCase       ,  past_key_values=_lowerCamelCase       ,  output_hidden_states=_lowerCamelCase       ,  )["""hidden_states"""][0]

         # select random slice
         __lowerCamelCase     :      List[str]       						=						ids_tensor((1,)       ,  output_from_past.shape[-1]							).item()
         __lowerCamelCase     :      str       						=						output_from_no_past[:, -3:, random_slice_idx].detach()
         __lowerCamelCase     :      List[Any]       						=						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(_lowerCamelCase       ,  _lowerCamelCase       ,  atol=1E-3							)							)


   def   _snake_case    (  self  : Any							):






         '''simple docstring'''



         __lowerCamelCase     :      Optional[Any]       						=						self.prepare_config_and_inputs()
         (
             (
             __lowerCamelCase
         )						, (
             __lowerCamelCase
         )						, (
             __lowerCamelCase
         )						, (
             __lowerCamelCase
         )						, (
             __lowerCamelCase
         )						, (
             __lowerCamelCase
         )						, (
             __lowerCamelCase
         )						, 
         )     :      str       						=						config_and_inputs
         __lowerCamelCase     :      str       						=						{"""input_ids""": input_ids, """attention_mask""": input_mask}
         return config, inputs_dict



@require_torch
class 							_UpperCamelCase (						A,A,A,unittest.TestCase						):

   '''simple docstring'''







   a_				:							Dict    = (
       (OpenLlamaModel, OpenLlamaForCausalLM, OpenLlamaForSequenceClassification) if is_torch_available() else ()
   )
   a_				:							Any    = (OpenLlamaForCausalLM,) if is_torch_available() else ()
   a_				:							Optional[int]    = (
       {
           "feature-extraction": OpenLlamaModel,
           "text-classification": OpenLlamaForSequenceClassification,
           "text-generation": OpenLlamaForCausalLM,
           "zero-shot": OpenLlamaForSequenceClassification,
       }
       if is_torch_available()
       else {}
   )
   a_				:							Any    = False
   a_				:							Dict    = False

   def   _snake_case    (  self  : int							):






         '''simple docstring'''



         __lowerCamelCase     :      Optional[int]       						=						OpenLlamaModelTester(self							)
         __lowerCamelCase     :      int       						=						ConfigTester(self       ,  config_class=_lowerCamelCase       ,  hidden_size=3_7							)

   def   _snake_case    (  self  : List[str]							):






         '''simple docstring'''



         self.config_tester.run_common_tests()

   def   _snake_case    (  self  : str							):






         '''simple docstring'''



         __lowerCamelCase     :      List[str]       						=						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*_lowerCamelCase							)

   def   _snake_case    (  self  : int							):






         '''simple docstring'''



         __lowerCamelCase     :      Union[str, Any]       						=						self.model_tester.prepare_config_and_inputs()
         for type in ["absolute", "relative_key", "relative_key_query"]:
               __lowerCamelCase     :      Optional[int]       						=						type
               self.model_tester.create_and_check_model(*_lowerCamelCase							)

   def   _snake_case    (  self  : Dict							):






         '''simple docstring'''



         __lowerCamelCase						, __lowerCamelCase     :      List[Any]       						=						self.model_tester.prepare_config_and_inputs_for_common()
         __lowerCamelCase     :      Tuple       						=						3
         __lowerCamelCase     :      int       						=						input_dict["""input_ids"""]
         __lowerCamelCase     :      List[str]       						=						input_ids.ne(1							).to(_lowerCamelCase							)
         __lowerCamelCase     :      Any       						=						ids_tensor([self.model_tester.batch_size]       ,  self.model_tester.type_sequence_label_size							)
         __lowerCamelCase     :      Tuple       						=						OpenLlamaForSequenceClassification(_lowerCamelCase							)
         model.to(_lowerCamelCase							)
         model.eval()
         __lowerCamelCase     :      Optional[int]       						=						model(_lowerCamelCase       ,  attention_mask=_lowerCamelCase       ,  labels=_lowerCamelCase							)
         self.assertEqual(result.logits.shape       ,  (self.model_tester.batch_size, self.model_tester.num_labels)							)

   def   _snake_case    (  self  : Any							):






         '''simple docstring'''



         __lowerCamelCase						, __lowerCamelCase     :      int       						=						self.model_tester.prepare_config_and_inputs_for_common()
         __lowerCamelCase     :      Optional[int]       						=						3
         __lowerCamelCase     :      Optional[Any]       						=						"""single_label_classification"""
         __lowerCamelCase     :      Any       						=						input_dict["""input_ids"""]
         __lowerCamelCase     :      Tuple       						=						input_ids.ne(1							).to(_lowerCamelCase							)
         __lowerCamelCase     :      int       						=						ids_tensor([self.model_tester.batch_size]       ,  self.model_tester.type_sequence_label_size							)
         __lowerCamelCase     :      List[str]       						=						OpenLlamaForSequenceClassification(_lowerCamelCase							)
         model.to(_lowerCamelCase							)
         model.eval()
         __lowerCamelCase     :      str       						=						model(_lowerCamelCase       ,  attention_mask=_lowerCamelCase       ,  labels=_lowerCamelCase							)
         self.assertEqual(result.logits.shape       ,  (self.model_tester.batch_size, self.model_tester.num_labels)							)

   def   _snake_case    (  self  : Any							):






         '''simple docstring'''



         __lowerCamelCase						, __lowerCamelCase     :      Tuple       						=						self.model_tester.prepare_config_and_inputs_for_common()
         __lowerCamelCase     :      Any       						=						3
         __lowerCamelCase     :      List[str]       						=						"""multi_label_classification"""
         __lowerCamelCase     :      Union[str, Any]       						=						input_dict["""input_ids"""]
         __lowerCamelCase     :      Optional[Any]       						=						input_ids.ne(1							).to(_lowerCamelCase							)
         __lowerCamelCase     :      Tuple       						=						ids_tensor(
             [self.model_tester.batch_size, config.num_labels]       ,  self.model_tester.type_sequence_label_size							).to(torch.float							)
         __lowerCamelCase     :      int       						=						OpenLlamaForSequenceClassification(_lowerCamelCase							)
         model.to(_lowerCamelCase							)
         model.eval()
         __lowerCamelCase     :      List[str]       						=						model(_lowerCamelCase       ,  attention_mask=_lowerCamelCase       ,  labels=_lowerCamelCase							)
         self.assertEqual(result.logits.shape       ,  (self.model_tester.batch_size, self.model_tester.num_labels)							)

   @unittest.skip("""Open-Llama buffers include complex numbers, which breaks this test"""							)
   def   _snake_case    (  self  : Dict							):






         '''simple docstring'''



         pass


   @parameterized.expand([("""linear""",), ("""dynamic""",)]							)
   def   _snake_case    (  self  : Any       ,  _lowerCamelCase  : Tuple							):






         '''simple docstring'''



         __lowerCamelCase						, __lowerCamelCase     :      List[Any]       						=						self.model_tester.prepare_config_and_inputs_for_common()
         __lowerCamelCase     :      Optional[Any]       						=						ids_tensor([1, 1_0]       ,  config.vocab_size							)
         __lowerCamelCase     :      Dict       						=						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
         __lowerCamelCase     :      str       						=						OpenLlamaModel(_lowerCamelCase							)
         original_model.to(_lowerCamelCase							)
         original_model.eval()
         __lowerCamelCase     :      Tuple       						=						original_model(_lowerCamelCase							).last_hidden_state
         __lowerCamelCase     :      Union[str, Any]       						=						original_model(_lowerCamelCase							).last_hidden_state

         set_seed(4_2							)  # Fixed seed at init time so the two models get the same random weights
         __lowerCamelCase     :      Dict       						=						{"""type""": scaling_type, """factor""": 10.0}
         __lowerCamelCase     :      int       						=						OpenLlamaModel(_lowerCamelCase							)
         scaled_model.to(_lowerCamelCase							)
         scaled_model.eval()
         __lowerCamelCase     :      List[Any]       						=						scaled_model(_lowerCamelCase							).last_hidden_state
         __lowerCamelCase     :      Union[str, Any]       						=						scaled_model(_lowerCamelCase							).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(_lowerCamelCase       ,  _lowerCamelCase       ,  atol=1E-5							)							)
         else:
               self.assertFalse(torch.allclose(_lowerCamelCase       ,  _lowerCamelCase       ,  atol=1E-5							)							)

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







from typing import Any





def   _UpperCAmelCase       (     UpperCAmelCase      :			list      ):
      """simple docstring"""




      if not input_list:
            return []
      __lowerCamelCase     :      int       						=						[input_list.count(UpperCAmelCase      ) for value in input_list]
      __lowerCamelCase     :      Tuple       						=						max(UpperCAmelCase      )  # Gets the maximum count in the input list.
      # Gets values of modes
      return sorted({input_list[i] for i, value in enumerate(UpperCAmelCase      ) if value == y}      )


if __name__ == "__main__":
    import doctest

    doctest.testmod()



