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 __snake_case				:




   def __init__(    self			:				List[Any]       ,	_lowercase			:				List[str]       ,	_lowercase			:				str=13       ,	_lowercase			:				Tuple=7       ,	_lowercase			:				List[str]=True       ,	_lowercase			:				Dict=True       ,	_lowercase			:				Union[str, Any]=False       ,	_lowercase			:				List[Any]=True       ,	_lowercase			:				str=99       ,	_lowercase			:				List[str]=32       ,	_lowercase			:				Optional[int]=5       ,	_lowercase			:				Dict=4       ,	_lowercase			:				Optional[Any]=37       ,	_lowercase			:				Tuple="gelu"       ,	_lowercase			:				Optional[Any]=0.1       ,	_lowercase			:				str=0.1       ,	_lowercase			:				Union[str, Any]=5_12       ,	_lowercase			:				List[Any]=16       ,	_lowercase			:				Dict=2       ,	_lowercase			:				Union[str, Any]=0.02       ,	_lowercase			:				Any=3       ,	_lowercase			:				Tuple=4       ,	_lowercase			:				Dict=None       ,	):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__										=      parent
     SCREAMING_SNAKE_CASE__										=      batch_size
     SCREAMING_SNAKE_CASE__										=      seq_length
     SCREAMING_SNAKE_CASE__										=      is_training
     SCREAMING_SNAKE_CASE__										=      use_input_mask
     SCREAMING_SNAKE_CASE__										=      use_token_type_ids
     SCREAMING_SNAKE_CASE__										=      use_labels
     SCREAMING_SNAKE_CASE__										=      vocab_size
     SCREAMING_SNAKE_CASE__										=      hidden_size
     SCREAMING_SNAKE_CASE__										=      num_hidden_layers
     SCREAMING_SNAKE_CASE__										=      num_attention_heads
     SCREAMING_SNAKE_CASE__										=      intermediate_size
     SCREAMING_SNAKE_CASE__										=      hidden_act
     SCREAMING_SNAKE_CASE__										=      hidden_dropout_prob
     SCREAMING_SNAKE_CASE__										=      attention_probs_dropout_prob
     SCREAMING_SNAKE_CASE__										=      max_position_embeddings
     SCREAMING_SNAKE_CASE__										=      type_vocab_size
     SCREAMING_SNAKE_CASE__										=      type_sequence_label_size
     SCREAMING_SNAKE_CASE__										=      initializer_range
     SCREAMING_SNAKE_CASE__										=      num_labels
     SCREAMING_SNAKE_CASE__										=      num_choices
     SCREAMING_SNAKE_CASE__										=      scope




   def 	__a						(    self			:				List[str]      ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__										=      ids_tensor([self.batch_size, self.seq_length]       ,	self.vocab_size      )

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

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

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

     SCREAMING_SNAKE_CASE__										=      self.get_config()

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




   def 	__a						(    self			:				List[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=A_       ,	initializer_range=self.initializer_range       ,	use_stable_embedding=A_       ,	)




   def 	__a						(    self			:				int       ,	_lowercase			:				Optional[int]       ,	_lowercase			:				Optional[int]       ,	_lowercase			:				Tuple       ,	_lowercase			:				Tuple       ,	_lowercase			:				List[str]       ,	_lowercase			:				Any       ,	_lowercase			:				int      ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__										=      OpenLlamaModel(config=A_      )
     model.to(A_      )
     model.eval()
     SCREAMING_SNAKE_CASE__										=      model(A_       ,	attention_mask=A_      )
     SCREAMING_SNAKE_CASE__										=      model(A_      )
     self.parent.assertEqual(result.last_hidden_state.shape       ,	(self.batch_size, self.seq_length, self.hidden_size)      )




   def 	__a						(    self			:				Any       ,	_lowercase			:				Optional[int]       ,	_lowercase			:				List[str]       ,	_lowercase			:				Dict       ,	_lowercase			:				List[Any]       ,	_lowercase			:				Optional[Any]       ,	_lowercase			:				List[Any]       ,	_lowercase			:				Optional[int]       ,	_lowercase			:				str       ,	_lowercase			:				Any       ,	):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__										=      True
     SCREAMING_SNAKE_CASE__										=      OpenLlamaModel(A_      )
     model.to(A_      )
     model.eval()
     SCREAMING_SNAKE_CASE__										=      model(
         A_       ,	attention_mask=A_       ,	encoder_hidden_states=A_       ,	encoder_attention_mask=A_       ,	)
     SCREAMING_SNAKE_CASE__										=      model(
         A_       ,	attention_mask=A_       ,	encoder_hidden_states=A_       ,	)
     SCREAMING_SNAKE_CASE__										=      model(A_       ,	attention_mask=A_      )
     self.parent.assertEqual(result.last_hidden_state.shape       ,	(self.batch_size, self.seq_length, self.hidden_size)      )




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



     """simple docstring"""






     SCREAMING_SNAKE_CASE__										=      OpenLlamaForCausalLM(config=A_      )
     model.to(A_      )
     model.eval()
     SCREAMING_SNAKE_CASE__										=      model(A_       ,	attention_mask=A_       ,	labels=A_      )
     self.parent.assertEqual(result.logits.shape       ,	(self.batch_size, self.seq_length, self.vocab_size)      )




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



     """simple docstring"""






     SCREAMING_SNAKE_CASE__										=      True
     SCREAMING_SNAKE_CASE__										=      True
     SCREAMING_SNAKE_CASE__										=      OpenLlamaForCausalLM(config=A_      )
     model.to(A_      )
     model.eval()

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

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

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

     SCREAMING_SNAKE_CASE__										=      model(
         A_       ,	attention_mask=A_       ,	encoder_hidden_states=A_       ,	encoder_attention_mask=A_       ,	output_hidden_states=A_       ,	)["""hidden_states"""][0]
     SCREAMING_SNAKE_CASE__										=      model(
         A_       ,	attention_mask=A_       ,	encoder_hidden_states=A_       ,	encoder_attention_mask=A_       ,	past_key_values=A_       ,	output_hidden_states=A_       ,	)["""hidden_states"""][0]

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




   def 	__a						(    self			:				List[str]      ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__										=      self.prepare_config_and_inputs()
     (
         (
         SCREAMING_SNAKE_CASE__      
     )							,     (
         SCREAMING_SNAKE_CASE__      
     )							,     (
         SCREAMING_SNAKE_CASE__      
     )							,     (
         SCREAMING_SNAKE_CASE__      
     )							,     (
         SCREAMING_SNAKE_CASE__      
     )							,     (
         SCREAMING_SNAKE_CASE__      
     )							,     (
         SCREAMING_SNAKE_CASE__      
     )							,           
     )										=      config_and_inputs
     SCREAMING_SNAKE_CASE__										=      {"""input_ids""": input_ids, """attention_mask""": input_mask}
     return config, inputs_dict







@require_torch
class __snake_case				(						A_   ,		A_   ,		A_   ,		unittest.TestCase		):
   lowerCAmelCase_														=							(
       (OpenLlamaModel, OpenLlamaForCausalLM, OpenLlamaForSequenceClassification) if is_torch_available() else ()
   )
   lowerCAmelCase_														=							(OpenLlamaForCausalLM,) if is_torch_available() else ()
   lowerCAmelCase_														=							(
       {
           "feature-extraction": OpenLlamaModel,
           "text-classification": OpenLlamaForSequenceClassification,
           "text-generation": OpenLlamaForCausalLM,
           "zero-shot": OpenLlamaForSequenceClassification,
       }
       if is_torch_available()
       else {}
   )
   lowerCAmelCase_														=							False
   lowerCAmelCase_														=							False




   def 	__a						(    self			:				Dict      ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__										=      OpenLlamaModelTester(self      )
     SCREAMING_SNAKE_CASE__										=      ConfigTester(self       ,	config_class=A_       ,	hidden_size=37      )




   def 	__a						(    self			:				Tuple      ):



     """simple docstring"""






     self.config_tester.run_common_tests()




   def 	__a						(    self			:				Any      ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__										=      self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*A_      )




   def 	__a						(    self			:				Optional[Any]      ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__										=      self.model_tester.prepare_config_and_inputs()
     for type in ["absolute", "relative_key", "relative_key_query"]:
       SCREAMING_SNAKE_CASE__										=      type
       self.model_tester.create_and_check_model(*A_      )




   def 	__a						(    self			:				Optional[Any]      ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__							,     SCREAMING_SNAKE_CASE__										=      self.model_tester.prepare_config_and_inputs_for_common()
     SCREAMING_SNAKE_CASE__										=      3
     SCREAMING_SNAKE_CASE__										=      input_dict["""input_ids"""]
     SCREAMING_SNAKE_CASE__										=      input_ids.ne(1      ).to(A_      )
     SCREAMING_SNAKE_CASE__										=      ids_tensor([self.model_tester.batch_size]       ,	self.model_tester.type_sequence_label_size      )
     SCREAMING_SNAKE_CASE__										=      OpenLlamaForSequenceClassification(A_      )
     model.to(A_      )
     model.eval()
     SCREAMING_SNAKE_CASE__										=      model(A_       ,	attention_mask=A_       ,	labels=A_      )
     self.assertEqual(result.logits.shape       ,	(self.model_tester.batch_size, self.model_tester.num_labels)      )




   def 	__a						(    self			:				Any      ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__							,     SCREAMING_SNAKE_CASE__										=      self.model_tester.prepare_config_and_inputs_for_common()
     SCREAMING_SNAKE_CASE__										=      3
     SCREAMING_SNAKE_CASE__										=      """single_label_classification"""
     SCREAMING_SNAKE_CASE__										=      input_dict["""input_ids"""]
     SCREAMING_SNAKE_CASE__										=      input_ids.ne(1      ).to(A_      )
     SCREAMING_SNAKE_CASE__										=      ids_tensor([self.model_tester.batch_size]       ,	self.model_tester.type_sequence_label_size      )
     SCREAMING_SNAKE_CASE__										=      OpenLlamaForSequenceClassification(A_      )
     model.to(A_      )
     model.eval()
     SCREAMING_SNAKE_CASE__										=      model(A_       ,	attention_mask=A_       ,	labels=A_      )
     self.assertEqual(result.logits.shape       ,	(self.model_tester.batch_size, self.model_tester.num_labels)      )




   def 	__a						(    self			:				Tuple      ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__							,     SCREAMING_SNAKE_CASE__										=      self.model_tester.prepare_config_and_inputs_for_common()
     SCREAMING_SNAKE_CASE__										=      3
     SCREAMING_SNAKE_CASE__										=      """multi_label_classification"""
     SCREAMING_SNAKE_CASE__										=      input_dict["""input_ids"""]
     SCREAMING_SNAKE_CASE__										=      input_ids.ne(1      ).to(A_      )
     SCREAMING_SNAKE_CASE__										=      ids_tensor(
         [self.model_tester.batch_size, config.num_labels]       ,	self.model_tester.type_sequence_label_size      ).to(torch.float      )
     SCREAMING_SNAKE_CASE__										=      OpenLlamaForSequenceClassification(A_      )
     model.to(A_      )
     model.eval()
     SCREAMING_SNAKE_CASE__										=      model(A_       ,	attention_mask=A_       ,	labels=A_      )
     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 	__a						(    self			:				Union[str, Any]      ):



     """simple docstring"""






     pass




   @parameterized.expand([("""linear""",), ("""dynamic""",)]      )
   def 	__a						(    self			:				str       ,	_lowercase			:				str      ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__							,     SCREAMING_SNAKE_CASE__										=      self.model_tester.prepare_config_and_inputs_for_common()
     SCREAMING_SNAKE_CASE__										=      ids_tensor([1, 10]       ,	config.vocab_size      )
     SCREAMING_SNAKE_CASE__										=      ids_tensor([1, int(config.max_position_embeddings * 1.5      )]       ,	config.vocab_size      )

     set_seed(42      )  # Fixed seed at init time so the two models get the same random weights
     SCREAMING_SNAKE_CASE__										=      OpenLlamaModel(A_      )
     original_model.to(A_      )
     original_model.eval()
     SCREAMING_SNAKE_CASE__										=      original_model(A_      ).last_hidden_state
     SCREAMING_SNAKE_CASE__										=      original_model(A_      ).last_hidden_state

     set_seed(42      )  # Fixed seed at init time so the two models get the same random weights
     SCREAMING_SNAKE_CASE__										=      {"""type""": scaling_type, """factor""": 10.0}
     SCREAMING_SNAKE_CASE__										=      OpenLlamaModel(A_      )
     scaled_model.to(A_      )
     scaled_model.eval()
     SCREAMING_SNAKE_CASE__										=      scaled_model(A_      ).last_hidden_state
     SCREAMING_SNAKE_CASE__										=      scaled_model(A_      ).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(A_       ,	A_       ,	atol=1E-5      )      )
     else:
       self.assertFalse(torch.allclose(A_       ,	A_       ,	atol=1E-5      )      )

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




from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_torch_available,
)


_A		       =						{'configuration_vit_mae': ['VIT_MAE_PRETRAINED_CONFIG_ARCHIVE_MAP', 'ViTMAEConfig']}

try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   _A		       =						[
       'VIT_MAE_PRETRAINED_MODEL_ARCHIVE_LIST',
       'ViTMAEForPreTraining',
       'ViTMAELayer',
       'ViTMAEModel',
       'ViTMAEPreTrainedModel',
   ]

try:
   if not is_tf_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   _A		       =						[
       'TFViTMAEForPreTraining',
       'TFViTMAEModel',
       'TFViTMAEPreTrainedModel',
   ]

if TYPE_CHECKING:
   from .configuration_vit_mae import VIT_MAE_PRETRAINED_CONFIG_ARCHIVE_MAP, ViTMAEConfig

   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_vit_mae import (
          VIT_MAE_PRETRAINED_MODEL_ARCHIVE_LIST,
          ViTMAEForPreTraining,
          ViTMAELayer,
          ViTMAEModel,
          ViTMAEPreTrainedModel,
      )

   try:
      if not is_tf_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_tf_vit_mae import TFViTMAEForPreTraining, TFViTMAEModel, TFViTMAEPreTrainedModel


else:
   import sys

   _A		       =						_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)





