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


       def __init__(    self						,       __A						,       __A=13						,       __A=7						,       __A=True						,       __A=True						,       __A=False						,       __A=True						,       __A=99						,       __A=32						,       __A=5						,       __A=4						,       __A=37						,       __A="gelu"						,       __A=0.1						,       __A=0.1						,       __A=512						,       __A=16						,       __A=2						,       __A=0.02						,       __A=3						,       __A=4						,       __A=None						,       )						->    Optional[int]:
              SCREAMING_SNAKE_CASE_							:				Any		=parent
              SCREAMING_SNAKE_CASE_							:				Optional[int]		=batch_size
              SCREAMING_SNAKE_CASE_							:				List[str]		=seq_length
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=is_training
              SCREAMING_SNAKE_CASE_							:				Tuple		=use_input_mask
              SCREAMING_SNAKE_CASE_							:				int		=use_token_type_ids
              SCREAMING_SNAKE_CASE_							:				List[str]		=use_labels
              SCREAMING_SNAKE_CASE_							:				Dict		=vocab_size
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=hidden_size
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=num_hidden_layers
              SCREAMING_SNAKE_CASE_							:				Optional[Any]		=num_attention_heads
              SCREAMING_SNAKE_CASE_							:				List[str]		=intermediate_size
              SCREAMING_SNAKE_CASE_							:				List[Any]		=hidden_act
              SCREAMING_SNAKE_CASE_							:				str		=hidden_dropout_prob
              SCREAMING_SNAKE_CASE_							:				Tuple		=attention_probs_dropout_prob
              SCREAMING_SNAKE_CASE_							:				Any		=max_position_embeddings
              SCREAMING_SNAKE_CASE_							:				Optional[int]		=type_vocab_size
              SCREAMING_SNAKE_CASE_							:				Any		=type_sequence_label_size
              SCREAMING_SNAKE_CASE_							:				Optional[int]		=initializer_range
              SCREAMING_SNAKE_CASE_							:				Any		=num_labels
              SCREAMING_SNAKE_CASE_							:				int		=num_choices
              SCREAMING_SNAKE_CASE_							:				Any		=scope


       def 	_snake_case				(    self				)						->    Dict:
              SCREAMING_SNAKE_CASE_							:				Optional[Any]		=ids_tensor([self.batch_size, self.seq_length]						,       self.vocab_size				)

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

              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=None
              if self.use_token_type_ids:
                     SCREAMING_SNAKE_CASE_							:				Optional[Any]		=ids_tensor([self.batch_size, self.seq_length]						,       self.type_vocab_size				)

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

              SCREAMING_SNAKE_CASE_							:				Optional[int]		=self.get_config()

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


       def 	_snake_case				(    self				)						->    Optional[Any]:
              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 	_snake_case				(    self						,       __A						,       __A						,       __A						,       __A						,       __A						,       __A						,       __A				)						->    Any:
              SCREAMING_SNAKE_CASE_							:				List[Any]		=OpenLlamaModel(config=__A				)
              model.to(__A				)
              model.eval()
              SCREAMING_SNAKE_CASE_							:				int		=model(__A						,       attention_mask=__A				)
              SCREAMING_SNAKE_CASE_							:				Tuple		=model(__A				)
              self.parent.assertEqual(result.last_hidden_state.shape						,       (self.batch_size, self.seq_length, self.hidden_size)				)


       def 	_snake_case				(    self						,       __A						,       __A						,       __A						,       __A						,       __A						,       __A						,       __A						,       __A						,       __A						,       )						->    Union[str, Any]:
              SCREAMING_SNAKE_CASE_							:				Dict		=True
              SCREAMING_SNAKE_CASE_							:				Optional[int]		=OpenLlamaModel(__A				)
              model.to(__A				)
              model.eval()
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=model(
                  __A						,       attention_mask=__A						,       encoder_hidden_states=__A						,       encoder_attention_mask=__A						,       )
              SCREAMING_SNAKE_CASE_							:				str		=model(
                  __A						,       attention_mask=__A						,       encoder_hidden_states=__A						,       )
              SCREAMING_SNAKE_CASE_							:				Any		=model(__A						,       attention_mask=__A				)
              self.parent.assertEqual(result.last_hidden_state.shape						,       (self.batch_size, self.seq_length, self.hidden_size)				)


       def 	_snake_case				(    self						,       __A						,       __A						,       __A						,       __A						,       __A						,       __A						,       __A						,       __A						,       __A						,       )						->    Optional[Any]:
              SCREAMING_SNAKE_CASE_							:				List[Any]		=OpenLlamaForCausalLM(config=__A				)
              model.to(__A				)
              model.eval()
              SCREAMING_SNAKE_CASE_							:				Tuple		=model(__A						,       attention_mask=__A						,       labels=__A				)
              self.parent.assertEqual(result.logits.shape						,       (self.batch_size, self.seq_length, self.vocab_size)				)


       def 	_snake_case				(    self						,       __A						,       __A						,       __A						,       __A						,       __A						,       __A						,       __A						,       __A						,       __A						,       )						->    Optional[int]:
              SCREAMING_SNAKE_CASE_							:				List[str]		=True
              SCREAMING_SNAKE_CASE_							:				Any		=True
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=OpenLlamaForCausalLM(config=__A				)
              model.to(__A				)
              model.eval()

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

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

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

              SCREAMING_SNAKE_CASE_							:				Any		=model(
                  __A						,       attention_mask=__A						,       encoder_hidden_states=__A						,       encoder_attention_mask=__A						,       output_hidden_states=__A						,       )['''hidden_states'''][0]
              SCREAMING_SNAKE_CASE_							:				List[Any]		=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_							:				List[Any]		=ids_tensor((1,)						,       output_from_past.shape[-1]				).item()
              SCREAMING_SNAKE_CASE_							:				Optional[Any]		=output_from_no_past[:, -3:, random_slice_idx].detach()
              SCREAMING_SNAKE_CASE_							:				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(__A						,       __A						,       atol=1e-3				)				)





       def 	_snake_case				(    self				)						->    List[Any]:
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=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_
              )	,     
              )							:				int		=config_and_inputs
              SCREAMING_SNAKE_CASE_							:				List[str]		={'''input_ids''': input_ids, '''attention_mask''': input_mask}
              return config, inputs_dict






@require_torch
class 			lowercase_		(      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 	_snake_case				(    self				)						->    List[Any]:
              SCREAMING_SNAKE_CASE_							:				str		=OpenLlamaModelTester(self				)
              SCREAMING_SNAKE_CASE_							:				List[Any]		=ConfigTester(self						,       config_class=__A						,       hidden_size=37				)


       def 	_snake_case				(    self				)						->    Tuple:
              self.config_tester.run_common_tests()


       def 	_snake_case				(    self				)						->    Tuple:
              SCREAMING_SNAKE_CASE_							:				str		=self.model_tester.prepare_config_and_inputs()
              self.model_tester.create_and_check_model(*__A				)


       def 	_snake_case				(    self				)						->    int:
              SCREAMING_SNAKE_CASE_							:				int		=self.model_tester.prepare_config_and_inputs()
              for type in ["absolute", "relative_key", "relative_key_query"]:
                     SCREAMING_SNAKE_CASE_							:				Dict		=type
                     self.model_tester.create_and_check_model(*__A				)


       def 	_snake_case				(    self				)						->    Any:
              SCREAMING_SNAKE_CASE_	,     SCREAMING_SNAKE_CASE_							:				Any		=self.model_tester.prepare_config_and_inputs_for_common()
              SCREAMING_SNAKE_CASE_							:				Dict		=3
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=input_dict['''input_ids''']
              SCREAMING_SNAKE_CASE_							:				Optional[int]		=input_ids.ne(1				).to(__A				)
              SCREAMING_SNAKE_CASE_							:				List[str]		=ids_tensor([self.model_tester.batch_size]						,       self.model_tester.type_sequence_label_size				)
              SCREAMING_SNAKE_CASE_							:				Any		=OpenLlamaForSequenceClassification(__A				)
              model.to(__A				)
              model.eval()
              SCREAMING_SNAKE_CASE_							:				List[str]		=model(__A						,       attention_mask=__A						,       labels=__A				)
              self.assertEqual(result.logits.shape						,       (self.model_tester.batch_size, self.model_tester.num_labels)				)


       def 	_snake_case				(    self				)						->    List[Any]:
              SCREAMING_SNAKE_CASE_	,     SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=self.model_tester.prepare_config_and_inputs_for_common()
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=3
              SCREAMING_SNAKE_CASE_							:				List[str]		='''single_label_classification'''
              SCREAMING_SNAKE_CASE_							:				Optional[int]		=input_dict['''input_ids''']
              SCREAMING_SNAKE_CASE_							:				int		=input_ids.ne(1				).to(__A				)
              SCREAMING_SNAKE_CASE_							:				str		=ids_tensor([self.model_tester.batch_size]						,       self.model_tester.type_sequence_label_size				)
              SCREAMING_SNAKE_CASE_							:				Optional[int]		=OpenLlamaForSequenceClassification(__A				)
              model.to(__A				)
              model.eval()
              SCREAMING_SNAKE_CASE_							:				int		=model(__A						,       attention_mask=__A						,       labels=__A				)
              self.assertEqual(result.logits.shape						,       (self.model_tester.batch_size, self.model_tester.num_labels)				)


       def 	_snake_case				(    self				)						->    Any:
              SCREAMING_SNAKE_CASE_	,     SCREAMING_SNAKE_CASE_							:				str		=self.model_tester.prepare_config_and_inputs_for_common()
              SCREAMING_SNAKE_CASE_							:				Dict		=3
              SCREAMING_SNAKE_CASE_							:				int		='''multi_label_classification'''
              SCREAMING_SNAKE_CASE_							:				str		=input_dict['''input_ids''']
              SCREAMING_SNAKE_CASE_							:				int		=input_ids.ne(1				).to(__A				)
              SCREAMING_SNAKE_CASE_							:				Tuple		=ids_tensor(
                  [self.model_tester.batch_size, config.num_labels]						,       self.model_tester.type_sequence_label_size				).to(torch.float				)
              SCREAMING_SNAKE_CASE_							:				Tuple		=OpenLlamaForSequenceClassification(__A				)
              model.to(__A				)
              model.eval()
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=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 	_snake_case				(    self				)						->    Dict:
              pass





       @parameterized.expand([('''linear''',), ('''dynamic''',)]				)
       def 	_snake_case				(    self						,       __A				)						->    str:
              SCREAMING_SNAKE_CASE_	,     SCREAMING_SNAKE_CASE_							:				Optional[Any]		=self.model_tester.prepare_config_and_inputs_for_common()
              SCREAMING_SNAKE_CASE_							:				Dict		=ids_tensor([1, 10]						,       config.vocab_size				)
              SCREAMING_SNAKE_CASE_							:				str		=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_							:				Any		=OpenLlamaModel(__A				)
              original_model.to(__A				)
              original_model.eval()
              SCREAMING_SNAKE_CASE_							:				List[Any]		=original_model(__A				).last_hidden_state
              SCREAMING_SNAKE_CASE_							:				Dict		=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_							:				Any		={'''type''': scaling_type, '''factor''': 10.0}
              SCREAMING_SNAKE_CASE_							:				Tuple		=OpenLlamaModel(__A				)
              scaled_model.to(__A				)
              scaled_model.eval()
              SCREAMING_SNAKE_CASE_							:				int		=scaled_model(__A				).last_hidden_state
              SCREAMING_SNAKE_CASE_							:				Tuple		=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_torch_available,
)


_lowercase											=    {
    """configuration_falcon""": ["""FALCON_PRETRAINED_CONFIG_ARCHIVE_MAP""", """FalconConfig"""],
}

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 _lowercase											=    [
     """FALCON_PRETRAINED_MODEL_ARCHIVE_LIST""",
     """FalconForCausalLM""",
     """FalconModel""",
     """FalconPreTrainedModel""",
     """FalconForSequenceClassification""",
     """FalconForTokenClassification""",
     """FalconForQuestionAnswering""",
 ]


if TYPE_CHECKING:
 from .configuration_falcon import FALCON_PRETRAINED_CONFIG_ARCHIVE_MAP, FalconConfig

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_falcon import (
      FALCON_PRETRAINED_MODEL_ARCHIVE_LIST,
      FalconForCausalLM,
      FalconForQuestionAnswering,
      FalconForSequenceClassification,
      FalconForTokenClassification,
      FalconModel,
      FalconPreTrainedModel,
  )


else:
 import sys

 _lowercase											=    _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)
