'''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      :			List[Any], UpperCamelCase__      :			Dict, UpperCamelCase__      :			Optional[int]=13, UpperCamelCase__      :			Tuple=7, UpperCamelCase__      :			Union[str, Any]=True, UpperCamelCase__      :			Any=True, UpperCamelCase__      :			int=False, UpperCamelCase__      :			str=True, UpperCamelCase__      :			Dict=99, UpperCamelCase__      :			Optional[Any]=32, UpperCamelCase__      :			Tuple=5, UpperCamelCase__      :			str=4, UpperCamelCase__      :			Optional[int]=37, UpperCamelCase__      :			List[str]="gelu", UpperCamelCase__      :			Optional[Any]=0.1, UpperCamelCase__      :			List[str]=0.1, UpperCamelCase__      :			Optional[int]=5_12, UpperCamelCase__      :			List[str]=16, UpperCamelCase__      :			List[Any]=2, UpperCamelCase__      :			Optional[int]=0.02, UpperCamelCase__      :			Dict=3, UpperCamelCase__      :			Optional[Any]=4, UpperCamelCase__      :			Tuple=None, )  ->			Dict:
        _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 	__UpperCAmelCase   (      self      :			int     )  ->			List[str]:
        _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 	__UpperCAmelCase   (      self      :			List[Any]     )  ->			Dict:
        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 	__UpperCAmelCase   (      self      :			Union[str, Any], UpperCamelCase__      :			int, UpperCamelCase__      :			Optional[Any], UpperCamelCase__      :			Optional[int], UpperCamelCase__      :			List[str], UpperCamelCase__      :			Dict, UpperCamelCase__      :			Optional[int], UpperCamelCase__      :			Any     )  ->			int:
        _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 	__UpperCAmelCase   (      self      :			List[str], UpperCamelCase__      :			Tuple, UpperCamelCase__      :			List[str], UpperCamelCase__      :			Tuple, UpperCamelCase__      :			str, UpperCamelCase__      :			int, UpperCamelCase__      :			Optional[Any], UpperCamelCase__      :			Union[str, Any], UpperCamelCase__      :			Union[str, Any], UpperCamelCase__      :			str, )  ->			List[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 	__UpperCAmelCase   (      self      :			Union[str, Any], UpperCamelCase__      :			Optional[int], UpperCamelCase__      :			Dict, UpperCamelCase__      :			Dict, UpperCamelCase__      :			Union[str, Any], UpperCamelCase__      :			Optional[Any], UpperCamelCase__      :			str, UpperCamelCase__      :			str, UpperCamelCase__      :			Dict, UpperCamelCase__      :			List[Any], )  ->			Any:
        _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 	__UpperCAmelCase   (      self      :			int, UpperCamelCase__      :			Any, UpperCamelCase__      :			int, UpperCamelCase__      :			List[Any], UpperCamelCase__      :			List[Any], UpperCamelCase__      :			Union[str, Any], UpperCamelCase__      :			List[Any], UpperCamelCase__      :			List[Any], UpperCamelCase__      :			Dict, UpperCamelCase__      :			int, )  ->			List[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 	__UpperCAmelCase   (      self      :			Optional[int]     )  ->			Any:
        _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_       (		_UpperCamelCase    ,    _UpperCamelCase    ,    _UpperCamelCase    ,    unittest.TestCase			):







   """simple docstring"""

   __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 	__UpperCAmelCase   (      self      :			Union[str, Any]     )  ->			Dict:
        _A    		=  OpenLlamaModelTester(self     )
        _A    		=  ConfigTester(self, config_class=UpperCamelCase__, hidden_size=37     )
   def 	__UpperCAmelCase   (      self      :			Dict     )  ->			Optional[int]:
        self.config_tester.run_common_tests()
   def 	__UpperCAmelCase   (      self      :			Optional[int]     )  ->			str:
        _A    		=  self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*UpperCamelCase__     )
   def 	__UpperCAmelCase   (      self      :			Optional[Any]     )  ->			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 	__UpperCAmelCase   (      self      :			int     )  ->			Union[str, 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 	__UpperCAmelCase   (      self      :			Union[str, Any]     )  ->			Any:
        _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 	__UpperCAmelCase   (      self      :			Union[str, Any]     )  ->			int:
        _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 	__UpperCAmelCase   (      self      :			Any     )  ->			List[str]:
        pass

   @parameterized.expand([('linear',), ('dynamic',)]     )
   def 	__UpperCAmelCase   (      self      :			Union[str, Any], UpperCamelCase__      :			Any     )  ->			List[str]:
        _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'''






import unittest

from diffusers.pipelines.pipeline_utils import is_safetensors_compatible







class lowercase_       (		unittest.TestCase			):







   """simple docstring"""

   def 	__UpperCAmelCase   (      self      :			Optional[Any]     )  ->			Dict:
        _A    		=  [
            'safety_checker/pytorch_model.bin',
            'safety_checker/model.safetensors',
            'vae/diffusion_pytorch_model.bin',
            'vae/diffusion_pytorch_model.safetensors',
            'text_encoder/pytorch_model.bin',
            'text_encoder/model.safetensors',
            'unet/diffusion_pytorch_model.bin',
            'unet/diffusion_pytorch_model.safetensors',
        ]
        self.assertTrue(is_safetensors_compatible(UpperCamelCase__     )     )
   def 	__UpperCAmelCase   (      self      :			Union[str, Any]     )  ->			Optional[Any]:
        _A    		=  [
            'unet/diffusion_pytorch_model.bin',
            'unet/diffusion_pytorch_model.safetensors',
        ]
        self.assertTrue(is_safetensors_compatible(UpperCamelCase__     )     )
   def 	__UpperCAmelCase   (      self      :			str     )  ->			Union[str, Any]:
        _A    		=  [
            'safety_checker/pytorch_model.bin',
            'safety_checker/model.safetensors',
            'vae/diffusion_pytorch_model.bin',
            'vae/diffusion_pytorch_model.safetensors',
            'text_encoder/pytorch_model.bin',
            'text_encoder/model.safetensors',
            'unet/diffusion_pytorch_model.bin',
            # Removed: 'unet/diffusion_pytorch_model.safetensors',
        ]
        self.assertFalse(is_safetensors_compatible(UpperCamelCase__     )     )
   def 	__UpperCAmelCase   (      self      :			int     )  ->			Dict:
        _A    		=  [
            'text_encoder/pytorch_model.bin',
            'text_encoder/model.safetensors',
        ]
        self.assertTrue(is_safetensors_compatible(UpperCamelCase__     )     )
   def 	__UpperCAmelCase   (      self      :			Union[str, Any]     )  ->			str:
        _A    		=  [
            'safety_checker/pytorch_model.bin',
            'safety_checker/model.safetensors',
            'vae/diffusion_pytorch_model.bin',
            'vae/diffusion_pytorch_model.safetensors',
            'text_encoder/pytorch_model.bin',
            # Removed: 'text_encoder/model.safetensors',
            'unet/diffusion_pytorch_model.bin',
            'unet/diffusion_pytorch_model.safetensors',
        ]
        self.assertFalse(is_safetensors_compatible(UpperCamelCase__     )     )
   def 	__UpperCAmelCase   (      self      :			str     )  ->			Dict:
        _A    		=  [
            'safety_checker/pytorch_model.fp16.bin',
            'safety_checker/model.fp16.safetensors',
            'vae/diffusion_pytorch_model.fp16.bin',
            'vae/diffusion_pytorch_model.fp16.safetensors',
            'text_encoder/pytorch_model.fp16.bin',
            'text_encoder/model.fp16.safetensors',
            'unet/diffusion_pytorch_model.fp16.bin',
            'unet/diffusion_pytorch_model.fp16.safetensors',
        ]
        _A    		=  'fp16'
        self.assertTrue(is_safetensors_compatible(UpperCamelCase__, variant=UpperCamelCase__     )     )
   def 	__UpperCAmelCase   (      self      :			Optional[int]     )  ->			List[str]:
        _A    		=  [
            'unet/diffusion_pytorch_model.fp16.bin',
            'unet/diffusion_pytorch_model.fp16.safetensors',
        ]
        _A    		=  'fp16'
        self.assertTrue(is_safetensors_compatible(UpperCamelCase__, variant=UpperCamelCase__     )     )
   def 	__UpperCAmelCase   (      self      :			List[str]     )  ->			List[str]:
        # pass variant but use the non-variant filenames
        _A    		=  [
            'unet/diffusion_pytorch_model.bin',
            'unet/diffusion_pytorch_model.safetensors',
        ]
        _A    		=  'fp16'
        self.assertTrue(is_safetensors_compatible(UpperCamelCase__, variant=UpperCamelCase__     )     )
   def 	__UpperCAmelCase   (      self      :			List[Any]     )  ->			Optional[Any]:
        _A    		=  [
            'safety_checker/pytorch_model.fp16.bin',
            'safety_checker/model.fp16.safetensors',
            'vae/diffusion_pytorch_model.fp16.bin',
            'vae/diffusion_pytorch_model.fp16.safetensors',
            'text_encoder/pytorch_model.fp16.bin',
            'text_encoder/model.fp16.safetensors',
            'unet/diffusion_pytorch_model.fp16.bin',
            # Removed: 'unet/diffusion_pytorch_model.fp16.safetensors',
        ]
        _A    		=  'fp16'
        self.assertFalse(is_safetensors_compatible(UpperCamelCase__, variant=UpperCamelCase__     )     )
   def 	__UpperCAmelCase   (      self      :			Dict     )  ->			Union[str, Any]:
        _A    		=  [
            'text_encoder/pytorch_model.fp16.bin',
            'text_encoder/model.fp16.safetensors',
        ]
        _A    		=  'fp16'
        self.assertTrue(is_safetensors_compatible(UpperCamelCase__, variant=UpperCamelCase__     )     )
   def 	__UpperCAmelCase   (      self      :			Tuple     )  ->			str:
        # pass variant but use the non-variant filenames
        _A    		=  [
            'text_encoder/pytorch_model.bin',
            'text_encoder/model.safetensors',
        ]
        _A    		=  'fp16'
        self.assertTrue(is_safetensors_compatible(UpperCamelCase__, variant=UpperCamelCase__     )     )

   def 	__UpperCAmelCase   (      self      :			List[Any]     )  ->			int:
        _A    		=  [
            'safety_checker/pytorch_model.fp16.bin',
            'safety_checker/model.fp16.safetensors',
            'vae/diffusion_pytorch_model.fp16.bin',
            'vae/diffusion_pytorch_model.fp16.safetensors',
            'text_encoder/pytorch_model.fp16.bin',
            # 'text_encoder/model.fp16.safetensors',
            'unet/diffusion_pytorch_model.fp16.bin',
            'unet/diffusion_pytorch_model.fp16.safetensors',
        ]
        _A    		=  'fp16'
        self.assertFalse(is_safetensors_compatible(UpperCamelCase__, variant=UpperCamelCase__     )     )



