





'''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       a_   :





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





        """simple docstring"""




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





     def 							lowercase__      ( self		:					List[Any]						):





        """simple docstring"""




        lowercase_			:Dict    			= ids_tensor([self.batch_size, self.seq_length]	,      self.vocab_size						)

        lowercase_			:List[str]    			= None
        if self.use_input_mask:
           lowercase_			:Optional[int]    			= random_attention_mask([self.batch_size, self.seq_length]						)

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

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

        lowercase_			:Any    			= self.get_config()

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





     def 							lowercase__      ( 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=lowercase	,      initializer_range=self.initializer_range	,      use_stable_embedding=lowercase	,      )





     def 							lowercase__      ( self		:					Optional[Any]	,      lowercase		:					Dict	,      lowercase		:					Optional[int]	,      lowercase		:					Optional[int]	,      lowercase		:					Any	,      lowercase		:					Optional[Any]	,      lowercase		:					int	,      lowercase		:					Dict						):





        """simple docstring"""




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





     def 							lowercase__      ( self		:					str	,      lowercase		:					Tuple	,      lowercase		:					str	,      lowercase		:					Union[str, Any]	,      lowercase		:					List[Any]	,      lowercase		:					List[Any]	,      lowercase		:					Optional[Any]	,      lowercase		:					Optional[int]	,      lowercase		:					List[Any]	,      lowercase		:					Union[str, Any]	,      ):





        """simple docstring"""




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





     def 							lowercase__      ( self		:					Dict	,      lowercase		:					Dict	,      lowercase		:					List[Any]	,      lowercase		:					Dict	,      lowercase		:					Dict	,      lowercase		:					List[str]	,      lowercase		:					Any	,      lowercase		:					Dict	,      lowercase		:					Union[str, Any]	,      lowercase		:					Any	,      ):





        """simple docstring"""




        lowercase_			:List[str]    			= OpenLlamaForCausalLM(config=lowercase						)
        model.to(lowercase						)
        model.eval()
        lowercase_			:List[str]    			= model(lowercase	,      attention_mask=lowercase	,      labels=lowercase						)
        self.parent.assertEqual(result.logits.shape	,      (self.batch_size, self.seq_length, self.vocab_size)						)





     def 							lowercase__      ( self		:					Optional[int]	,      lowercase		:					str	,      lowercase		:					Tuple	,      lowercase		:					List[Any]	,      lowercase		:					Any	,      lowercase		:					Any	,      lowercase		:					Tuple	,      lowercase		:					Optional[Any]	,      lowercase		:					str	,      lowercase		:					Tuple	,      ):





        """simple docstring"""




        lowercase_			:int    			= True
        lowercase_			:str    			= True
        lowercase_			:Union[str, Any]    			= OpenLlamaForCausalLM(config=lowercase						)
        model.to(lowercase						)
        model.eval()

        # first forward pass
        lowercase_			:Optional[Any]    			= model(
            lowercase	,      attention_mask=lowercase	,      encoder_hidden_states=lowercase	,      encoder_attention_mask=lowercase	,      use_cache=lowercase	,      )
        lowercase_			:Any    			= outputs.past_key_values

        # create hypothetical multiple next token and extent to next_input_ids
        lowercase_			:List[str]    			= ids_tensor((self.batch_size, 3)	,      config.vocab_size						)
        lowercase_			:Dict    			= ids_tensor((self.batch_size, 3)	,      vocab_size=2						)

        # append to next input_ids and
        lowercase_			:Union[str, Any]    			= torch.cat([input_ids, next_tokens]	,      dim=-1						)
        lowercase_			:Union[str, Any]    			= torch.cat([input_mask, next_mask]	,      dim=-1						)

        lowercase_			:List[Any]    			= model(
            lowercase	,      attention_mask=lowercase	,      encoder_hidden_states=lowercase	,      encoder_attention_mask=lowercase	,      output_hidden_states=lowercase	,      )["hidden_states"][0]
        lowercase_			:str    			= model(
            lowercase	,      attention_mask=lowercase	,      encoder_hidden_states=lowercase	,      encoder_attention_mask=lowercase	,      past_key_values=lowercase	,      output_hidden_states=lowercase	,      )["hidden_states"][0]

        # select random slice
        lowercase_			:int    			= ids_tensor((1,)	,      output_from_past.shape[-1]						).item()
        lowercase_			:Union[str, Any]    			= output_from_no_past[:, -3:, random_slice_idx].detach()
        lowercase_			:Optional[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(lowercase	,      lowercase	,      atol=1e-3						)						)







     def 							lowercase__      ( self		:					int						):





        """simple docstring"""




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







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





     def 							lowercase__      ( self		:					str						):





        """simple docstring"""




        lowercase_			:Union[str, Any]    			= OpenLlamaModelTester(self						)
        lowercase_			:int    			= ConfigTester(self	,      config_class=lowercase	,      hidden_size=37						)





     def 							lowercase__      ( self		:					Optional[int]						):





        """simple docstring"""




        self.config_tester.run_common_tests()





     def 							lowercase__      ( self		:					str						):





        """simple docstring"""




        lowercase_			:Any    			= self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*lowercase						)





     def 							lowercase__      ( self		:					List[str]						):





        """simple docstring"""




        lowercase_			:Any    			= self.model_tester.prepare_config_and_inputs()
        for type in ["absolute", "relative_key", "relative_key_query"]:
           lowercase_			:List[Any]    			= type
           self.model_tester.create_and_check_model(*lowercase						)





     def 							lowercase__      ( self		:					Dict						):





        """simple docstring"""




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





     def 							lowercase__      ( self		:					Dict						):





        """simple docstring"""




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





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





        """simple docstring"""




        lowercase_			,							lowercase_			:Any    			= self.model_tester.prepare_config_and_inputs_for_common()
        lowercase_			:Optional[Any]    			= 3
        lowercase_			:Dict    			= "multi_label_classification"
        lowercase_			:Optional[int]    			= input_dict["input_ids"]
        lowercase_			:List[str]    			= input_ids.ne(1						).to(lowercase						)
        lowercase_			:Dict    			= ids_tensor(
            [self.model_tester.batch_size, config.num_labels]	,      self.model_tester.type_sequence_label_size						).to(torch.float						)
        lowercase_			:Tuple    			= OpenLlamaForSequenceClassification(lowercase						)
        model.to(lowercase						)
        model.eval()
        lowercase_			:Optional[Any]    			= model(lowercase	,      attention_mask=lowercase	,      labels=lowercase						)
        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 							lowercase__      ( self		:					Optional[Any]						):





        """simple docstring"""




        pass







     @parameterized.expand([("linear",), ("dynamic",)]						)
     def 							lowercase__      ( self		:					int	,      lowercase		:					Tuple						):





        """simple docstring"""




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

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

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









'''simple docstring'''


from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


lowerCAmelCase					:     int												={
    '''configuration_bigbird_pegasus''': [
        '''BIGBIRD_PEGASUS_PRETRAINED_CONFIG_ARCHIVE_MAP''',
        '''BigBirdPegasusConfig''',
        '''BigBirdPegasusOnnxConfig''',
    ],
}

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 lowerCAmelCase					:     Any												=[
     '''BIGBIRD_PEGASUS_PRETRAINED_MODEL_ARCHIVE_LIST''',
     '''BigBirdPegasusForCausalLM''',
     '''BigBirdPegasusForConditionalGeneration''',
     '''BigBirdPegasusForQuestionAnswering''',
     '''BigBirdPegasusForSequenceClassification''',
     '''BigBirdPegasusModel''',
     '''BigBirdPegasusPreTrainedModel''',
 ]


if TYPE_CHECKING:
 from .configuration_bigbird_pegasus import (
     BIGBIRD_PEGASUS_PRETRAINED_CONFIG_ARCHIVE_MAP,
     BigBirdPegasusConfig,
     BigBirdPegasusOnnxConfig,
 )

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_bigbird_pegasus import (
      BIGBIRD_PEGASUS_PRETRAINED_MODEL_ARCHIVE_LIST,
      BigBirdPegasusForCausalLM,
      BigBirdPegasusForConditionalGeneration,
      BigBirdPegasusForQuestionAnswering,
      BigBirdPegasusForSequenceClassification,
      BigBirdPegasusModel,
      BigBirdPegasusPreTrainedModel,
  )


else:
 import sys

 lowerCAmelCase					:     Dict												=_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)


