




"""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 							lowercase_			:



 '''simple docstring'''

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

 def      lowerCAmelCase_				(					self			:			Any						):
       _A     =  ids_tensor([self.batch_size, self.seq_length]       ,				self.vocab_size						)

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

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

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

       _A     =  self.get_config()

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

 def      lowerCAmelCase_				(					self			:			str						):
       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			:			Tuple       ,				_UpperCAmelCase			:			Any       ,				_UpperCAmelCase			:			Union[str, Any]       ,				_UpperCAmelCase			:			Optional[Any]       ,				_UpperCAmelCase			:			str       ,				_UpperCAmelCase			:			Union[str, Any]       ,				_UpperCAmelCase			:			Dict       ,				_UpperCAmelCase			:			Any						):
       _A     =  OpenLlamaModel(config=_UpperCAmelCase						)
       model.to(_UpperCAmelCase						)
       model.eval()
       _A     =  model(_UpperCAmelCase       ,				attention_mask=_UpperCAmelCase						)
       _A     =  model(_UpperCAmelCase						)
       self.parent.assertEqual(result.last_hidden_state.shape       ,				(self.batch_size, self.seq_length, self.hidden_size)						)

 def      lowerCAmelCase_				(					self			:			Any       ,				_UpperCAmelCase			:			int       ,				_UpperCAmelCase			:			Optional[int]       ,				_UpperCAmelCase			:			Tuple       ,				_UpperCAmelCase			:			Dict       ,				_UpperCAmelCase			:			Optional[Any]       ,				_UpperCAmelCase			:			Optional[int]       ,				_UpperCAmelCase			:			List[Any]       ,				_UpperCAmelCase			:			Tuple       ,				_UpperCAmelCase			:			Optional[Any]       ,				):
       _A     =  True
       _A     =  OpenLlamaModel(_UpperCAmelCase						)
       model.to(_UpperCAmelCase						)
       model.eval()
       _A     =  model(
           _UpperCAmelCase       ,				attention_mask=_UpperCAmelCase       ,				encoder_hidden_states=_UpperCAmelCase       ,				encoder_attention_mask=_UpperCAmelCase       ,				)
       _A     =  model(
           _UpperCAmelCase       ,				attention_mask=_UpperCAmelCase       ,				encoder_hidden_states=_UpperCAmelCase       ,				)
       _A     =  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			:			Tuple       ,				_UpperCAmelCase			:			Optional[Any]       ,				_UpperCAmelCase			:			Union[str, Any]       ,				_UpperCAmelCase			:			List[str]       ,				_UpperCAmelCase			:			str       ,				_UpperCAmelCase			:			int       ,				_UpperCAmelCase			:			Union[str, Any]       ,				_UpperCAmelCase			:			int       ,				_UpperCAmelCase			:			str       ,				_UpperCAmelCase			:			Optional[int]       ,				):
       _A     =  OpenLlamaForCausalLM(config=_UpperCAmelCase						)
       model.to(_UpperCAmelCase						)
       model.eval()
       _A     =  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			:			Tuple       ,				_UpperCAmelCase			:			Union[str, Any]       ,				_UpperCAmelCase			:			List[Any]       ,				_UpperCAmelCase			:			Dict       ,				_UpperCAmelCase			:			Optional[int]       ,				_UpperCAmelCase			:			Optional[int]       ,				_UpperCAmelCase			:			List[str]       ,				_UpperCAmelCase			:			Any       ,				_UpperCAmelCase			:			Optional[Any]       ,				_UpperCAmelCase			:			Any       ,				):
       _A     =  True
       _A     =  True
       _A     =  OpenLlamaForCausalLM(config=_UpperCAmelCase						)
       model.to(_UpperCAmelCase						)
       model.eval()

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

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

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

       _A     =  model(
           _UpperCAmelCase       ,				attention_mask=_UpperCAmelCase       ,				encoder_hidden_states=_UpperCAmelCase       ,				encoder_attention_mask=_UpperCAmelCase       ,				output_hidden_states=_UpperCAmelCase       ,				)['hidden_states'][0]
       _A     =  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
       _A     =  ids_tensor((1,)       ,				output_from_past.shape[-1]						).item()
       _A     =  output_from_no_past[:, -3:, random_slice_idx].detach()
       _A     =  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			:			List[str]						):
       _A     =  self.prepare_config_and_inputs()
       (
           (
           _A
       )  ,       (
           _A
       )  ,       (
           _A
       )  ,       (
           _A
       )  ,       (
           _A
       )  ,       (
           _A
       )  ,       (
           _A
       )  ,       
       )     =  config_and_inputs
       _A     =  {'input_ids': input_ids, 'attention_mask': input_mask}
       return config, inputs_dict





@require_torch
class 							lowercase_			(	__lowerCAmelCase     ,      __lowerCAmelCase     ,      __lowerCAmelCase     ,      unittest.TestCase				):



 '''simple docstring'''
 UpperCAmelCase       :	List[str]     =      (
     (OpenLlamaModel, OpenLlamaForCausalLM, OpenLlamaForSequenceClassification) if is_torch_available() else ()
 )
 UpperCAmelCase       :	List[str]     =      (OpenLlamaForCausalLM,) if is_torch_available() else ()
 UpperCAmelCase       :	str     =      (
     {
         '''feature-extraction''': OpenLlamaModel,
         '''text-classification''': OpenLlamaForSequenceClassification,
         '''text-generation''': OpenLlamaForCausalLM,
         '''zero-shot''': OpenLlamaForSequenceClassification,
     }
     if is_torch_available()
     else {}
 )
 UpperCAmelCase       :	Any     =      False
 UpperCAmelCase       :	Any     =      False

 def      lowerCAmelCase_				(					self			:			Union[str, Any]						):
       _A     =  OpenLlamaModelTester(self						)
       _A     =  ConfigTester(self       ,				config_class=_UpperCAmelCase       ,				hidden_size=37						)

 def      lowerCAmelCase_				(					self			:			List[Any]						):
       self.config_tester.run_common_tests()

 def      lowerCAmelCase_				(					self			:			Optional[Any]						):
       _A     =  self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*_UpperCAmelCase						)

 def      lowerCAmelCase_				(					self			:			Tuple						):
       _A     =  self.model_tester.prepare_config_and_inputs()
       for type in ["absolute", "relative_key", "relative_key_query"]:
             _A     =  type
             self.model_tester.create_and_check_model(*_UpperCAmelCase						)

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



 @parameterized.expand([('linear',), ('dynamic',)]						)
 def      lowerCAmelCase_				(					self			:			Optional[int]       ,				_UpperCAmelCase			:			Union[str, Any]						):
       _A  ,       _A     =  self.model_tester.prepare_config_and_inputs_for_common()
       _A     =  ids_tensor([1, 10]       ,				config.vocab_size						)
       _A     =  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
       _A     =  OpenLlamaModel(_UpperCAmelCase						)
       original_model.to(_UpperCAmelCase						)
       original_model.eval()
       _A     =  original_model(_UpperCAmelCase						).last_hidden_state
       _A     =  original_model(_UpperCAmelCase						).last_hidden_state

       set_seed(42						)  # Fixed seed at init time so the two models get the same random weights
       _A     =  {'type': scaling_type, 'factor': 10.0}
       _A     =  OpenLlamaModel(_UpperCAmelCase						)
       scaled_model.to(_UpperCAmelCase						)
       scaled_model.eval()
       _A     =  scaled_model(_UpperCAmelCase						).last_hidden_state
       _A     =  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						)						)







"""simple docstring"""

from .configuration_bert_masked import MaskedBertConfig
from .modeling_bert_masked import (
    MaskedBertForMultipleChoice,
    MaskedBertForQuestionAnswering,
    MaskedBertForSequenceClassification,
    MaskedBertForTokenClassification,
    MaskedBertModel,
)
from .modules import *

