



"""simple docstring"""






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    :



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



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



   def    __UpperCAmelCase    (       self    :	List[str]				):



     """simple docstring"""
     _snake_case   	=      ids_tensor([self.batch_size, self.seq_length]							,       self.vocab_size				)

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

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

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

     _snake_case   	=      self.get_config()

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



   def    __UpperCAmelCase    (       self    :	str				):



     """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=SCREAMING_SNAKE_CASE__							,       initializer_range=self.initializer_range							,       use_stable_embedding=SCREAMING_SNAKE_CASE__							,       )



   def    __UpperCAmelCase    (       self    :	Union[str, Any]							,       __lowerCamelCase    :	str							,       __lowerCamelCase    :	Any							,       __lowerCamelCase    :	int							,       __lowerCamelCase    :	List[str]							,       __lowerCamelCase    :	str							,       __lowerCamelCase    :	List[str]							,       __lowerCamelCase    :	List[str]				):



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



   def    __UpperCAmelCase    (       self    :	List[str]							,       __lowerCamelCase    :	List[str]							,       __lowerCamelCase    :	str							,       __lowerCamelCase    :	List[str]							,       __lowerCamelCase    :	Optional[Any]							,       __lowerCamelCase    :	int							,       __lowerCamelCase    :	List[str]							,       __lowerCamelCase    :	Any							,       __lowerCamelCase    :	Dict							,       __lowerCamelCase    :	Any							,       ):



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



   def    __UpperCAmelCase    (       self    :	str							,       __lowerCamelCase    :	List[str]							,       __lowerCamelCase    :	Optional[int]							,       __lowerCamelCase    :	str							,       __lowerCamelCase    :	Dict							,       __lowerCamelCase    :	List[Any]							,       __lowerCamelCase    :	Dict							,       __lowerCamelCase    :	Optional[int]							,       __lowerCamelCase    :	str							,       __lowerCamelCase    :	Union[str, Any]							,       ):



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



   def    __UpperCAmelCase    (       self    :	int							,       __lowerCamelCase    :	List[str]							,       __lowerCamelCase    :	int							,       __lowerCamelCase    :	Tuple							,       __lowerCamelCase    :	Tuple							,       __lowerCamelCase    :	Union[str, Any]							,       __lowerCamelCase    :	int							,       __lowerCamelCase    :	List[str]							,       __lowerCamelCase    :	List[Any]							,       __lowerCamelCase    :	Union[str, Any]							,       ):



     """simple docstring"""
     _snake_case   	=      True
     _snake_case   	=      True
     _snake_case   	=      OpenLlamaForCausalLM(config=SCREAMING_SNAKE_CASE__				)
     model.to(SCREAMING_SNAKE_CASE__				)
     model.eval()

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

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

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

     _snake_case   	=      model(
         SCREAMING_SNAKE_CASE__							,       attention_mask=SCREAMING_SNAKE_CASE__							,       encoder_hidden_states=SCREAMING_SNAKE_CASE__							,       encoder_attention_mask=SCREAMING_SNAKE_CASE__							,       output_hidden_states=SCREAMING_SNAKE_CASE__							,       )['''hidden_states'''][0]
     _snake_case   	=      model(
         SCREAMING_SNAKE_CASE__							,       attention_mask=SCREAMING_SNAKE_CASE__							,       encoder_hidden_states=SCREAMING_SNAKE_CASE__							,       encoder_attention_mask=SCREAMING_SNAKE_CASE__							,       past_key_values=SCREAMING_SNAKE_CASE__							,       output_hidden_states=SCREAMING_SNAKE_CASE__							,       )['''hidden_states'''][0]

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





   def    __UpperCAmelCase    (       self    :	str				):



     """simple docstring"""
     _snake_case   	=      self.prepare_config_and_inputs()
     (
         (
         _snake_case		
     )	,     (
         _snake_case		
     )	,     (
         _snake_case		
     )	,     (
         _snake_case		
     )	,     (
         _snake_case		
     )	,     (
         _snake_case		
     )	,     (
         _snake_case		
     )	,     		
     )   	=      config_and_inputs
     _snake_case   	=      {'''input_ids''': input_ids, '''attention_mask''': input_mask}
     return config, inputs_dict



@require_torch
class     UpperCAmelCase    (	__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE,unittest.TestCase  ):
   A__		:	Union[str, Any]           =  (
       (OpenLlamaModel, OpenLlamaForCausalLM, OpenLlamaForSequenceClassification) if is_torch_available() else ()
   )
   A__		:	Dict           =  (OpenLlamaForCausalLM,) if is_torch_available() else ()
   A__		:	str           =  (
       {
           "feature-extraction": OpenLlamaModel,
           "text-classification": OpenLlamaForSequenceClassification,
           "text-generation": OpenLlamaForCausalLM,
           "zero-shot": OpenLlamaForSequenceClassification,
       }
       if is_torch_available()
       else {}
   )
   A__		:	Optional[Any]           =  False
   A__		:	Any           =  False



   def    __UpperCAmelCase    (       self    :	int				):



     """simple docstring"""
     _snake_case   	=      OpenLlamaModelTester(self				)
     _snake_case   	=      ConfigTester(self							,       config_class=SCREAMING_SNAKE_CASE__							,       hidden_size=3_7				)



   def    __UpperCAmelCase    (       self    :	str				):



     """simple docstring"""
     self.config_tester.run_common_tests()



   def    __UpperCAmelCase    (       self    :	List[Any]				):



     """simple docstring"""
     _snake_case   	=      self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE__				)



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



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



   def    __UpperCAmelCase    (       self    :	Optional[Any]				):



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



   def    __UpperCAmelCase    (       self    :	List[str]				):



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



   def    __UpperCAmelCase    (       self    :	Dict				):



     """simple docstring"""
     _snake_case	,     _snake_case   	=      self.model_tester.prepare_config_and_inputs_for_common()
     _snake_case   	=      3
     _snake_case   	=      '''multi_label_classification'''
     _snake_case   	=      input_dict['''input_ids''']
     _snake_case   	=      input_ids.ne(1				).to(SCREAMING_SNAKE_CASE__				)
     _snake_case   	=      ids_tensor(
         [self.model_tester.batch_size, config.num_labels]							,       self.model_tester.type_sequence_label_size				).to(torch.float				)
     _snake_case   	=      OpenLlamaForSequenceClassification(SCREAMING_SNAKE_CASE__				)
     model.to(SCREAMING_SNAKE_CASE__				)
     model.eval()
     _snake_case   	=      model(SCREAMING_SNAKE_CASE__							,       attention_mask=SCREAMING_SNAKE_CASE__							,       labels=SCREAMING_SNAKE_CASE__				)
     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    __UpperCAmelCase    (       self    :	Tuple				):



     """simple docstring"""
     pass





   @parameterized.expand([('''linear''',), ('''dynamic''',)]				)
   def    __UpperCAmelCase    (       self    :	Tuple							,       __lowerCamelCase    :	Any				):



     """simple docstring"""
     _snake_case	,     _snake_case   	=      self.model_tester.prepare_config_and_inputs_for_common()
     _snake_case   	=      ids_tensor([1, 1_0]							,       config.vocab_size				)
     _snake_case   	=      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
     _snake_case   	=      OpenLlamaModel(SCREAMING_SNAKE_CASE__				)
     original_model.to(SCREAMING_SNAKE_CASE__				)
     original_model.eval()
     _snake_case   	=      original_model(SCREAMING_SNAKE_CASE__				).last_hidden_state
     _snake_case   	=      original_model(SCREAMING_SNAKE_CASE__				).last_hidden_state

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

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







import logging

from transformers import PretrainedConfig


_snake_case      = logging.getLogger(__name__)


_snake_case      = {
    '''bertabs-finetuned-cnndm''': '''https://huggingface.co/remi/bertabs-finetuned-cnndm-extractive-abstractive-summarization/resolve/main/config.json''',
}
class       _lowerCAmelCase  (		__magic_name__							):

 """simple docstring"""
 SCREAMING_SNAKE_CASE_				:					int       						="bertabs"

 def __init__(						self			:			Optional[int]					, SCREAMING_SNAKE_CASE__			:			Union[str, Any]=3_05_22					, SCREAMING_SNAKE_CASE__			:			Union[str, Any]=5_12					, SCREAMING_SNAKE_CASE__			:			List[Any]=6					, SCREAMING_SNAKE_CASE__			:			List[str]=5_12					, SCREAMING_SNAKE_CASE__			:			Union[str, Any]=8					, SCREAMING_SNAKE_CASE__			:			List[str]=5_12					, SCREAMING_SNAKE_CASE__			:			List[Any]=0.2					, SCREAMING_SNAKE_CASE__			:			Tuple=6					, SCREAMING_SNAKE_CASE__			:			Any=7_68					, SCREAMING_SNAKE_CASE__			:			str=8					, SCREAMING_SNAKE_CASE__			:			Dict=20_48					, SCREAMING_SNAKE_CASE__			:			List[str]=0.2					, **SCREAMING_SNAKE_CASE__			:			List[str]					, ):






      """simple docstring"""



      super().__init__(**SCREAMING_SNAKE_CASE__   )

      UpperCamelCase						    =					vocab_size
      UpperCamelCase						    =					max_pos

      UpperCamelCase						    =					enc_layers
      UpperCamelCase						    =					enc_hidden_size
      UpperCamelCase						    =					enc_heads
      UpperCamelCase						    =					enc_ff_size
      UpperCamelCase						    =					enc_dropout

      UpperCamelCase						    =					dec_layers
      UpperCamelCase						    =					dec_hidden_size
      UpperCamelCase						    =					dec_heads
      UpperCamelCase						    =					dec_ff_size
      UpperCamelCase						    =					dec_dropout


