


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


       """simple docstring"""




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


             '''simple docstring'''



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





       def 						lowerCAmelCase     (   self							:			Optional[Any]       ):


             '''simple docstring'''



             lowerCamelCase__		:				Dict    =  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__		:				Any    =  ids_tensor([self.batch_size, self.seq_length]					,      self.type_vocab_size       )

             lowerCamelCase__		:				Tuple    =  None
             lowerCamelCase__		:				int    =  None
             lowerCamelCase__		:				Optional[int]    =  None
             if self.use_labels:
                   lowerCamelCase__		:				str    =  ids_tensor([self.batch_size]					,      self.type_sequence_label_size       )
                   lowerCamelCase__		:				Optional[int]    =  ids_tensor([self.batch_size, self.seq_length]					,      self.num_labels       )
                   lowerCamelCase__		:				str    =  ids_tensor([self.batch_size]					,      self.num_choices       )

             lowerCamelCase__		:				Any    =  self.get_config()

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





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





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


             '''simple docstring'''



             lowerCamelCase__		:				Optional[Any]    =  OpenLlamaModel(config=UpperCamelCase__       )
             model.to(UpperCamelCase__       )
             model.eval()
             lowerCamelCase__		:				Optional[int]    =  model(UpperCamelCase__					,      attention_mask=UpperCamelCase__       )
             lowerCamelCase__		:				List[str]    =  model(UpperCamelCase__       )
             self.parent.assertEqual(result.last_hidden_state.shape					,      (self.batch_size, self.seq_length, self.hidden_size)       )





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


             '''simple docstring'''



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





       def 						lowerCAmelCase     (   self							:			Optional[int]					,      __lowerCamelCase							:			Tuple					,      __lowerCamelCase							:			Any					,      __lowerCamelCase							:			Union[str, Any]					,      __lowerCamelCase							:			Optional[Any]					,      __lowerCamelCase							:			Any					,      __lowerCamelCase							:			Dict					,      __lowerCamelCase							:			Dict					,      __lowerCamelCase							:			Any					,      __lowerCamelCase							:			Tuple					,      ):


             '''simple docstring'''



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





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


             '''simple docstring'''



             lowerCamelCase__		:				Optional[Any]    =  True
             lowerCamelCase__		:				Optional[int]    =  True
             lowerCamelCase__		:				Any    =  OpenLlamaForCausalLM(config=UpperCamelCase__       )
             model.to(UpperCamelCase__       )
             model.eval()

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

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

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

             lowerCamelCase__		:				Any    =  model(
                 UpperCamelCase__					,      attention_mask=UpperCamelCase__					,      encoder_hidden_states=UpperCamelCase__					,      encoder_attention_mask=UpperCamelCase__					,      output_hidden_states=UpperCamelCase__					,      )["hidden_states"][0]
             lowerCamelCase__		:				Optional[int]    =  model(
                 UpperCamelCase__					,      attention_mask=UpperCamelCase__					,      encoder_hidden_states=UpperCamelCase__					,      encoder_attention_mask=UpperCamelCase__					,      past_key_values=UpperCamelCase__					,      output_hidden_states=UpperCamelCase__					,      )["hidden_states"][0]

             # select random slice
             lowerCamelCase__		:				int    =  ids_tensor((1,)					,      output_from_past.shape[-1]       ).item()
             lowerCamelCase__		:				int    =  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(UpperCamelCase__					,      UpperCamelCase__					,      atol=1E-3       )       )





       def 						lowerCAmelCase     (   self							:			Optional[Any]       ):


             '''simple docstring'''



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







@require_torch
class 			_lowercase (							_A							,  _A							,  _A							,  unittest.TestCase):


       """simple docstring"""




       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 						lowerCAmelCase     (   self							:			Union[str, Any]       ):


             '''simple docstring'''



             lowerCamelCase__		:				Any    =  OpenLlamaModelTester(self       )
             lowerCamelCase__		:				List[Any]    =  ConfigTester(self					,      config_class=UpperCamelCase__					,      hidden_size=37       )





       def 						lowerCAmelCase     (   self							:			List[Any]       ):


             '''simple docstring'''



             self.config_tester.run_common_tests()





       def 						lowerCAmelCase     (   self							:			str       ):


             '''simple docstring'''



             lowerCamelCase__		:				Any    =  self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_model(*UpperCamelCase__       )





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


             '''simple docstring'''



             lowerCamelCase__		:				Optional[Any]    =  self.model_tester.prepare_config_and_inputs()
             for type in ["absolute", "relative_key", "relative_key_query"]:
                   lowerCamelCase__		:				str    =  type
                   self.model_tester.create_and_check_model(*UpperCamelCase__       )





       def 						lowerCAmelCase     (   self							:			Optional[Any]       ):


             '''simple docstring'''



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





       def 						lowerCAmelCase     (   self							:			List[str]       ):


             '''simple docstring'''



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





       def 						lowerCAmelCase     (   self							:			str       ):


             '''simple docstring'''



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


             '''simple docstring'''



             pass





       @parameterized.expand([("linear",), ("dynamic",)]       )
       def 						lowerCAmelCase     (   self							:			Tuple					,      __lowerCamelCase							:			List[Any]       ):


             '''simple docstring'''



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

             set_seed(42       )  # Fixed seed at init time so the two models get the same random weights
             lowerCamelCase__		:				int    =  {"type": scaling_type, "factor": 1_0.0}
             lowerCamelCase__		:				Tuple    =  OpenLlamaModel(UpperCamelCase__       )
             scaled_model.to(UpperCamelCase__       )
             scaled_model.eval()
             lowerCamelCase__		:				List[str]    =  scaled_model(UpperCamelCase__       ).last_hidden_state
             lowerCamelCase__		:				str    =  scaled_model(UpperCamelCase__       ).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(UpperCamelCase__					,      UpperCamelCase__					,      atol=1E-5       )       )
             else:
                   self.assertFalse(torch.allclose(UpperCamelCase__					,      UpperCamelCase__					,      atol=1E-5       )       )

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




from typing import TYPE_CHECKING

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


__lowerCAmelCase			:     str           =   {
    'configuration_funnel': ['FUNNEL_PRETRAINED_CONFIG_ARCHIVE_MAP', 'FunnelConfig'],
    'convert_funnel_original_tf_checkpoint_to_pytorch': [],
    'tokenization_funnel': ['FunnelTokenizer'],
}

try:
    if not is_tokenizers_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    __lowerCAmelCase			:     Any           =   ['FunnelTokenizerFast']

try:
    if not is_torch_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    __lowerCAmelCase			:     Optional[int]           =   [
        'FUNNEL_PRETRAINED_MODEL_ARCHIVE_LIST',
        'FunnelBaseModel',
        'FunnelForMaskedLM',
        'FunnelForMultipleChoice',
        'FunnelForPreTraining',
        'FunnelForQuestionAnswering',
        'FunnelForSequenceClassification',
        'FunnelForTokenClassification',
        'FunnelModel',
        'FunnelPreTrainedModel',
        'load_tf_weights_in_funnel',
    ]

try:
    if not is_tf_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    __lowerCAmelCase			:     Tuple           =   [
        'TF_FUNNEL_PRETRAINED_MODEL_ARCHIVE_LIST',
        'TFFunnelBaseModel',
        'TFFunnelForMaskedLM',
        'TFFunnelForMultipleChoice',
        'TFFunnelForPreTraining',
        'TFFunnelForQuestionAnswering',
        'TFFunnelForSequenceClassification',
        'TFFunnelForTokenClassification',
        'TFFunnelModel',
        'TFFunnelPreTrainedModel',
    ]


if TYPE_CHECKING:
    from .configuration_funnel import FUNNEL_PRETRAINED_CONFIG_ARCHIVE_MAP, FunnelConfig
    from .tokenization_funnel import FunnelTokenizer

    try:
        if not is_tokenizers_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .tokenization_funnel_fast import FunnelTokenizerFast

    try:
        if not is_torch_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .modeling_funnel import (
            FUNNEL_PRETRAINED_MODEL_ARCHIVE_LIST,
            FunnelBaseModel,
            FunnelForMaskedLM,
            FunnelForMultipleChoice,
            FunnelForPreTraining,
            FunnelForQuestionAnswering,
            FunnelForSequenceClassification,
            FunnelForTokenClassification,
            FunnelModel,
            FunnelPreTrainedModel,
            load_tf_weights_in_funnel,
        )

    try:
        if not is_tf_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .modeling_tf_funnel import (
            TF_FUNNEL_PRETRAINED_MODEL_ARCHIVE_LIST,
            TFFunnelBaseModel,
            TFFunnelForMaskedLM,
            TFFunnelForMultipleChoice,
            TFFunnelForPreTraining,
            TFFunnelForQuestionAnswering,
            TFFunnelForSequenceClassification,
            TFFunnelForTokenClassification,
            TFFunnelModel,
            TFFunnelPreTrainedModel,
        )

else:
    import sys

    __lowerCAmelCase			:     List[Any]           =   _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)

