'''simple docstring'''

import unittest

from parameterized import parameterized

from transformers import LlamaConfig, is_torch_available, set_seed
from transformers.testing_utils import require_torch, slow, 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 LlamaForCausalLM, LlamaForSequenceClassification, LlamaModel, LlamaTokenizer



class       UpperCAmelCase  :





       '''simple docstring'''



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

       def _lowerCAmelCase(						self			)  ->  Dict:
        lowercase__    :   Optional[int]						= ids_tensor([self.batch_size, self.seq_length]     ,       self.vocab_size			)

        lowercase__    :   str						= None
        if self.use_input_mask:
         lowercase__    :   List[Any]						= random_attention_mask([self.batch_size, self.seq_length]			)

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

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

        lowercase__    :   List[Any]						= self.get_config()

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

       def _lowerCAmelCase(						self			)  ->  List[Any]:
        return LlamaConfig(
            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=__lowerCAmelCase     ,       initializer_range=self.initializer_range     ,       )

       def _lowerCAmelCase(						self     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase			)  ->  List[Any]:
        lowercase__    :   int						= LlamaModel(config=__lowerCAmelCase			)
        model.to(__lowerCAmelCase			)
        model.eval()
        lowercase__    :   Optional[Any]						= model(__lowerCAmelCase     ,       attention_mask=__lowerCAmelCase			)
        lowercase__    :   Optional[Any]						= model(__lowerCAmelCase			)
        self.parent.assertEqual(result.last_hidden_state.shape     ,       (self.batch_size, self.seq_length, self.hidden_size)			)

       def _lowerCAmelCase(						self     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       )  ->  Optional[int]:
        lowercase__    :   Tuple						= True
        lowercase__    :   Optional[Any]						= LlamaModel(__lowerCAmelCase			)
        model.to(__lowerCAmelCase			)
        model.eval()
        lowercase__    :   Any						= model(
            __lowerCAmelCase     ,       attention_mask=__lowerCAmelCase     ,       encoder_hidden_states=__lowerCAmelCase     ,       encoder_attention_mask=__lowerCAmelCase     ,       )
        lowercase__    :   Any						= model(
            __lowerCAmelCase     ,       attention_mask=__lowerCAmelCase     ,       encoder_hidden_states=__lowerCAmelCase     ,       )
        lowercase__    :   Tuple						= model(__lowerCAmelCase     ,       attention_mask=__lowerCAmelCase			)
        self.parent.assertEqual(result.last_hidden_state.shape     ,       (self.batch_size, self.seq_length, self.hidden_size)			)

       def _lowerCAmelCase(						self     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       )  ->  List[str]:
        lowercase__    :   Optional[int]						= LlamaForCausalLM(config=__lowerCAmelCase			)
        model.to(__lowerCAmelCase			)
        model.eval()
        lowercase__    :   Union[str, Any]						= model(__lowerCAmelCase     ,       attention_mask=__lowerCAmelCase     ,       labels=__lowerCAmelCase			)
        self.parent.assertEqual(result.logits.shape     ,       (self.batch_size, self.seq_length, self.vocab_size)			)

       def _lowerCAmelCase(						self     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       )  ->  str:
        lowercase__    :   str						= True
        lowercase__    :   Tuple						= True
        lowercase__    :   Optional[Any]						= LlamaForCausalLM(config=__lowerCAmelCase			)
        model.to(__lowerCAmelCase			)
        model.eval()

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

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

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

        lowercase__    :   Any						= model(
            __lowerCAmelCase     ,       attention_mask=__lowerCAmelCase     ,       encoder_hidden_states=__lowerCAmelCase     ,       encoder_attention_mask=__lowerCAmelCase     ,       output_hidden_states=__lowerCAmelCase     ,       )['''hidden_states'''][0]
        lowercase__    :   Dict						= model(
            __lowerCAmelCase     ,       attention_mask=__lowerCAmelCase     ,       encoder_hidden_states=__lowerCAmelCase     ,       encoder_attention_mask=__lowerCAmelCase     ,       past_key_values=__lowerCAmelCase     ,       output_hidden_states=__lowerCAmelCase     ,       )['''hidden_states'''][0]

        # select random slice
        lowercase__    :   List[str]						= ids_tensor((1,)     ,       output_from_past.shape[-1]			).item()
        lowercase__    :   str						= output_from_no_past[:, -3:, random_slice_idx].detach()
        lowercase__    :   int						= 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(__lowerCAmelCase     ,       __lowerCAmelCase     ,       atol=1E-3			)			)


       def _lowerCAmelCase(						self			)  ->  Dict:
        lowercase__    :   Tuple						= self.prepare_config_and_inputs()
        (
            (
            lowercase__
        )			,							(
            lowercase__
        )			,							(
            lowercase__
        )			,							(
            lowercase__
        )			,							(
            lowercase__
        )			,							(
            lowercase__
        )			,							(
            lowercase__
        )			,							
        )    :   Union[str, Any]						= config_and_inputs
        lowercase__    :   Dict						= {'''input_ids''': input_ids, '''attention_mask''': input_mask}
        return config, inputs_dict



@require_torch
class       UpperCAmelCase  (						a__   ,						a__   ,						a__   ,						unittest.TestCase ):





       '''simple docstring'''



       SCREAMING_SNAKE_CASE 		=   (LlamaModel, LlamaForCausalLM, LlamaForSequenceClassification) if is_torch_available() else ()
       SCREAMING_SNAKE_CASE 		=   (LlamaForCausalLM,) if is_torch_available() else ()
       SCREAMING_SNAKE_CASE 		=   (
           {
               "feature-extraction": LlamaModel,
               "text-classification": LlamaForSequenceClassification,
               "text-generation": LlamaForCausalLM,
               "zero-shot": LlamaForSequenceClassification,
           }
           if is_torch_available()
           else {}
       )
       SCREAMING_SNAKE_CASE 		=   False
       SCREAMING_SNAKE_CASE 		=   False

       def _lowerCAmelCase(						self			)  ->  Union[str, Any]:
        lowercase__    :   str						= LlamaModelTester(self			)
        lowercase__    :   Optional[int]						= ConfigTester(self     ,       config_class=__lowerCAmelCase     ,       hidden_size=37			)

       def _lowerCAmelCase(						self			)  ->  int:
        self.config_tester.run_common_tests()

       def _lowerCAmelCase(						self			)  ->  int:
        lowercase__    :   Tuple						= self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*__lowerCAmelCase			)

       def _lowerCAmelCase(						self			)  ->  List[Any]:
        lowercase__    :   Union[str, Any]						= self.model_tester.prepare_config_and_inputs()
        for type in ["absolute", "relative_key", "relative_key_query"]:
         lowercase__    :   List[str]						= type
         self.model_tester.create_and_check_model(*__lowerCAmelCase			)

       def _lowerCAmelCase(						self			)  ->  Optional[int]:
        lowercase__			,							lowercase__    :   Optional[int]						= self.model_tester.prepare_config_and_inputs_for_common()
        lowercase__    :   Dict						= 3
        lowercase__    :   Tuple						= input_dict['''input_ids''']
        lowercase__    :   List[Any]						= input_ids.ne(1			).to(__lowerCAmelCase			)
        lowercase__    :   List[str]						= ids_tensor([self.model_tester.batch_size]     ,       self.model_tester.type_sequence_label_size			)
        lowercase__    :   Optional[Any]						= LlamaForSequenceClassification(__lowerCAmelCase			)
        model.to(__lowerCAmelCase			)
        model.eval()
        lowercase__    :   Optional[Any]						= model(__lowerCAmelCase     ,       attention_mask=__lowerCAmelCase     ,       labels=__lowerCAmelCase			)
        self.assertEqual(result.logits.shape     ,       (self.model_tester.batch_size, self.model_tester.num_labels)			)

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

       def _lowerCAmelCase(						self			)  ->  List[str]:
        lowercase__			,							lowercase__    :   Optional[Any]						= self.model_tester.prepare_config_and_inputs_for_common()
        lowercase__    :   Dict						= 3
        lowercase__    :   int						= '''multi_label_classification'''
        lowercase__    :   List[str]						= input_dict['''input_ids''']
        lowercase__    :   Optional[int]						= input_ids.ne(1			).to(__lowerCAmelCase			)
        lowercase__    :   Dict						= ids_tensor(
            [self.model_tester.batch_size, config.num_labels]     ,       self.model_tester.type_sequence_label_size			).to(torch.float			)
        lowercase__    :   int						= LlamaForSequenceClassification(__lowerCAmelCase			)
        model.to(__lowerCAmelCase			)
        model.eval()
        lowercase__    :   Any						= model(__lowerCAmelCase     ,       attention_mask=__lowerCAmelCase     ,       labels=__lowerCAmelCase			)
        self.assertEqual(result.logits.shape     ,       (self.model_tester.batch_size, self.model_tester.num_labels)			)

       @unittest.skip('''LLaMA buffers include complex numbers, which breaks this test'''			)
       def _lowerCAmelCase(						self			)  ->  Any:
        pass


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

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

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



@require_torch
class       UpperCAmelCase  (						unittest.TestCase ):





       '''simple docstring'''



       @unittest.skip('''Logits are not exactly the same, once we fix the instabalities somehow, will update!'''			)
       @slow
       def _lowerCAmelCase(						self			)  ->  Union[str, Any]:
        lowercase__    :   Any						= [1, 306, 4658, 278, 6593, 310, 2834, 338]
        lowercase__    :   List[Any]						= LlamaForCausalLM.from_pretrained('''meta-llama/Llama-2-7b-hf'''     ,       device_map='''auto'''			)
        lowercase__    :   Optional[Any]						= model(torch.tensor([input_ids]			)			)
        # Expected mean on dim = -1
        lowercase__    :   Tuple						= torch.tensor([[-6.6_5_5_0, -4.1_2_2_7, -4.9_8_5_9, -3.2_4_0_6, 0.8_2_6_2, -3.0_0_3_3, 1.2_9_6_4, -3.3_6_9_9]]			)
        torch.testing.assert_close(out.mean(-1			)     ,       __lowerCAmelCase     ,       atol=1E-2     ,       rtol=1E-2			)
        # slicing logits[0, 0, 0:30]
        # fmt: off
        lowercase__    :   Optional[Any]						= torch.tensor([-1_2.8_2_8_1, -7.4_4_5_3, -0.4_6_3_9, -8.0_6_2_5, -7.2_5_0_0, -8.0_0_0_0, -6.4_8_8_3, -7.7_6_9_5, -7.8_4_3_8, -7.0_3_1_2, -6.2_1_8_8, -7.1_3_2_8, -1.8_4_9_6, 1.9_9_6_1, -8.6_2_5_0, -6.7_2_2_7, -1_2.8_2_8_1, -6.9_4_9_2, -7.0_7_4_2, -7.7_8_5_2, -7.5_8_2_0, -7.9_0_6_2, -6.9_3_7_5, -7.9_8_0_5, -8.3_4_3_8, -8.1_5_6_2, -8.0_4_6_9, -7.6_2_5_0, -7.7_4_2_2, -7.3_3_9_8,]			)
        # fmt: on
        torch.testing.assert_close(out[0, 0, :30]     ,       __lowerCAmelCase     ,       atol=1E-5     ,       rtol=1E-5			)

       @unittest.skip('''Logits are not exactly the same, once we fix the instabalities somehow, will update!'''			)
       @slow
       def _lowerCAmelCase(						self			)  ->  Tuple:
        lowercase__    :   Optional[Any]						= [1, 306, 4658, 278, 6593, 310, 2834, 338]
        lowercase__    :   Tuple						= LlamaForCausalLM.from_pretrained('''meta-llama/Llama-2-13b-hf'''     ,       device_map='''auto'''			)
        lowercase__    :   Dict						= model(torch.tensor(__lowerCAmelCase			)			)
        # Expected mean on dim = -1
        lowercase__    :   Dict						= torch.tensor([[-2.0_6_2_2, -1.2_7_9_4, -1.1_6_3_8, -0.9_7_8_8, -1.4_6_0_3, -1.0_2_3_8, -1.7_8_9_3, -1.4_4_1_1]]			)
        torch.testing.assert_close(out.mean(-1			)     ,       __lowerCAmelCase     ,       atol=1E-2     ,       rtol=1E-2			)
        # slicing logits[0, 0, 0:30]
        # fmt: off
        lowercase__    :   Dict						= torch.tensor([-8.1_4_0_6, -8.0_5_4_7, 2.7_4_6_1, -1.2_3_4_4, -0.1_4_4_8, -1.8_2_6_2, -1.0_0_2_0, -1.8_1_5_4, -1.6_8_9_5, -1.8_5_1_6, -2.3_5_7_4, -0.9_2_7_7, 3.7_5_9_8, 6.5_7_4_2, -1.2_9_9_8, -0.1_1_7_7, -8.1_4_0_6, -2.9_6_8_8, -2.9_1_9_9, -3.1_6_9_9, -3.5_2_5_4, -2.3_5_5_5, -2.7_9_8_8, -3.4_1_4_1, -2.8_2_6_2, -4.5_1_9_5, -3.3_3_7_9, -3.3_1_6_4, -2.7_8_3_2, -3.0_2_7_3]			)
        # fmt: on
        torch.testing.assert_close(out[0, 0, :30]     ,       __lowerCAmelCase     ,       atol=1E-5     ,       rtol=1E-5			)

       @unittest.skip('''Logits are not exactly the same, once we fix the instabalities somehow, will update!'''			)
       @slow
       def _lowerCAmelCase(						self			)  ->  Any:
        lowercase__    :   Any						= [1, 306, 4658, 278, 6593, 310, 2834, 338]
        lowercase__    :   List[Any]						= LlamaForCausalLM.from_pretrained('''meta-llama/Llama-2-13b-chat-hf'''     ,       device_map='''auto'''			)
        lowercase__    :   Dict						= model(torch.tensor(__lowerCAmelCase			)			)
        # Expected mean on dim = -1
        lowercase__    :   Optional[Any]						= torch.tensor([[-0.8_5_6_2, -1.8_5_2_0, -0.7_5_5_1, -0.4_1_6_2, -1.5_1_6_1, -1.2_0_3_8, -2.4_8_2_3, -2.3_2_5_4]]			)
        torch.testing.assert_close(out.mean(-1			)     ,       __lowerCAmelCase     ,       atol=1E-2     ,       rtol=1E-2			)
        # slicing logits[0, 0, 0:30]
        # fmt: off
        lowercase__    :   Any						= torch.tensor([-2.2_2_2_7, 4.8_8_2_8, 0.9_0_2_3, -0.4_5_7_8, -0.7_8_7_1, -0.1_0_3_3, -0.6_2_2_1, -0.5_7_8_6, -0.7_8_0_3, -1.0_6_7_4, -1.2_9_2_0, -0.1_5_7_0, 0.8_0_0_8, 2.0_7_2_3, -0.9_4_9_7, 0.2_7_7_1, -2.2_2_2_7, -0.7_6_1_2, -1.4_3_4_6, -1.2_0_6_1, -1.6_4_2_6, -0.3_0_0_0, -0.7_1_3_9, -1.1_9_3_4, -1.8_6_9_1, -1.6_9_7_3, -1.5_9_4_7, -1.2_7_0_5, -0.3_5_2_3, -0.5_5_1_3]			)
        # fmt: on
        torch.testing.assert_close(out.mean(-1			)     ,       __lowerCAmelCase     ,       atol=1E-2     ,       rtol=1E-2			)

       @unittest.skip(
           '''Logits are not exactly the same, once we fix the instabalities somehow, will update! Also it is gonna be a `too_slow` test'''			)
       @slow
       def _lowerCAmelCase(						self			)  ->  List[Any]:
        lowercase__    :   Any						= [1, 306, 4658, 278, 6593, 310, 2834, 338]
        lowercase__    :   Optional[int]						= LlamaForCausalLM.from_pretrained('''meta-llama/Llama-2-70b-hf'''     ,       device_map='''auto'''			)
        lowercase__    :   Dict						= model(torch.tensor(__lowerCAmelCase			)			)

        lowercase__    :   Optional[Any]						= torch.tensor(
            [[-4.2_3_2_7, -3.3_3_6_0, -4.6_6_6_5, -4.7_6_3_1, -1.8_1_8_0, -3.4_1_7_0, -1.4_2_1_1, -3.1_8_1_0]]     ,       dtype=torch.floataa			)
        torch.testing.assert_close(out.mean(-1			)     ,       __lowerCAmelCase     ,       atol=1E-2     ,       rtol=1E-2			)
        # fmt: off
        lowercase__    :   Any						= torch.tensor([-9.4_9_2_2, -3.9_5_5_1, 1.7_9_9_8, -5.6_7_5_8, -5.1_0_5_5, -5.8_9_8_4, -4.8_3_2_0, -6.8_0_8_6, -6.5_3_9_1, -5.6_1_7_2, -5.5_8_2_0, -5.5_3_5_2, 1.7_8_8_1, 3.6_2_8_9, -6.5_1_1_7, -3.4_7_8_5, -9.5_0_0_0, -6.0_3_5_2, -6.8_1_2_5, -6.0_1_9_5, -6.6_8_3_6, -5.4_7_2_7, -6.2_8_1_2, -6.0_3_9_1, -7.3_3_9_8, -7.4_2_9_7, -7.4_8_4_4, -6.5_8_2_0, -5.8_7_8_9, -5.5_3_1_2]			)
        # fmt: on
        torch.testing.assert_close(out[0, 0, :30]     ,       __lowerCAmelCase     ,       atol=1E-5     ,       rtol=1E-5			)


       @unittest.skip('''Model is curently gated'''			)
       @slow
       def _lowerCAmelCase(						self			)  ->  Tuple:
        lowercase__    :   Optional[int]						= '''Simply put, the theory of relativity states that 1) the laws of physics are the same everywhere in the universe and 2) the passage of time and the length of objects can vary depending on the observer\'s frame of reference.\n\nThe first part of the theory, that the laws of physics are the same everywhere, is known as the "princi'''
        lowercase__    :   Dict						= '''Simply put, the theory of relativity states that '''
        lowercase__    :   List[str]						= LlamaTokenizer.from_pretrained('''meta-llama/Llama-2-13b-chat-hf'''			)
        lowercase__    :   str						= tokenizer.encode(__lowerCAmelCase     ,       return_tensors='''pt'''			)
        lowercase__    :   str						= LlamaForCausalLM.from_pretrained(
            '''meta-llama/Llama-2-13b-chat-hf'''     ,       device_map='''sequential'''     ,       use_safetensors=__lowerCAmelCase			)

        # greedy generation outputs
        lowercase__    :   Dict						= model.generate(__lowerCAmelCase     ,       max_new_tokens=64     ,       top_p=__lowerCAmelCase     ,       temperature=1     ,       do_sample=__lowerCAmelCase			)
        lowercase__    :   List[Any]						= tokenizer.decode(generated_ids[0]     ,       skip_special_tokens=__lowerCAmelCase			)
        self.assertEqual(__lowerCAmelCase     ,       __lowerCAmelCase			)


'''simple docstring'''

import unittest

import numpy as np
import torch

from diffusers import DDIMPipeline, DDIMScheduler, UNetaDModel
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, slow, torch_device

from ..pipeline_params import UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS, UNCONDITIONAL_IMAGE_GENERATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()



class       UpperCAmelCase  (						a__   ,						unittest.TestCase ):





       '''simple docstring'''



       SCREAMING_SNAKE_CASE 		=   DDIMPipeline
       SCREAMING_SNAKE_CASE 		=   UNCONDITIONAL_IMAGE_GENERATION_PARAMS
       SCREAMING_SNAKE_CASE 		=   PipelineTesterMixin.required_optional_params - {
           "num_images_per_prompt",
           "latents",
           "callback",
           "callback_steps",
       }
       SCREAMING_SNAKE_CASE 		=   UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS
       SCREAMING_SNAKE_CASE 		=   False

       def _lowerCAmelCase(						self			)  ->  Union[str, Any]:
        torch.manual_seed(0			)
        lowercase__    :   Optional[int]						= UNetaDModel(
            block_out_channels=(32, 64)     ,       layers_per_block=2     ,       sample_size=32     ,       in_channels=3     ,       out_channels=3     ,       down_block_types=('''DownBlock2D''', '''AttnDownBlock2D''')     ,       up_block_types=('''AttnUpBlock2D''', '''UpBlock2D''')     ,       )
        lowercase__    :   Any						= DDIMScheduler()
        lowercase__    :   Union[str, Any]						= {'''unet''': unet, '''scheduler''': scheduler}
        return components

       def _lowerCAmelCase(						self     ,       __lowerCAmelCase     ,       __lowerCAmelCase=0			)  ->  int:
        if str(__lowerCAmelCase			).startswith('''mps'''			):
         lowercase__    :   Optional[Any]						= torch.manual_seed(__lowerCAmelCase			)
        else:
         lowercase__    :   int						= torch.Generator(device=__lowerCAmelCase			).manual_seed(__lowerCAmelCase			)
        lowercase__    :   List[Any]						= {
            '''batch_size''': 1,
            '''generator''': generator,
            '''num_inference_steps''': 2,
            '''output_type''': '''numpy''',
        }
        return inputs

       def _lowerCAmelCase(						self			)  ->  Optional[Any]:
        lowercase__    :   int						= '''cpu'''

        lowercase__    :   Optional[Any]						= self.get_dummy_components()
        lowercase__    :   str						= self.pipeline_class(**__lowerCAmelCase			)
        pipe.to(__lowerCAmelCase			)
        pipe.set_progress_bar_config(disable=__lowerCAmelCase			)

        lowercase__    :   str						= self.get_dummy_inputs(__lowerCAmelCase			)
        lowercase__    :   str						= pipe(**__lowerCAmelCase			).images
        lowercase__    :   List[str]						= image[0, -3:, -3:, -1]

        self.assertEqual(image.shape     ,       (1, 32, 32, 3)			)
        lowercase__    :   Optional[int]						= np.array(
            [1.000E00, 5.717E-01, 4.717E-01, 1.000E00, 0.000E00, 1.000E00, 3.000E-04, 0.000E00, 9.000E-04]			)
        lowercase__    :   Optional[Any]						= np.abs(image_slice.flatten() - expected_slice			).max()
        self.assertLessEqual(__lowerCAmelCase     ,       1E-3			)

       def _lowerCAmelCase(						self			)  ->  str:
        super().test_dict_tuple_outputs_equivalent(expected_max_difference=3E-3			)

       def _lowerCAmelCase(						self			)  ->  List[str]:
        super().test_save_load_local(expected_max_difference=3E-3			)

       def _lowerCAmelCase(						self			)  ->  Tuple:
        super().test_save_load_optional_components(expected_max_difference=3E-3			)


       def _lowerCAmelCase(						self			)  ->  int:
        super().test_inference_batch_single_identical(expected_max_diff=3E-3			)



@slow
@require_torch_gpu
class       UpperCAmelCase  (						unittest.TestCase ):





       '''simple docstring'''



       def _lowerCAmelCase(						self			)  ->  List[Any]:
        lowercase__    :   List[str]						= '''google/ddpm-cifar10-32'''

        lowercase__    :   List[Any]						= UNetaDModel.from_pretrained(__lowerCAmelCase			)
        lowercase__    :   str						= DDIMScheduler()

        lowercase__    :   List[Any]						= DDIMPipeline(unet=__lowerCAmelCase     ,       scheduler=__lowerCAmelCase			)
        ddim.to(__lowerCAmelCase			)
        ddim.set_progress_bar_config(disable=__lowerCAmelCase			)

        lowercase__    :   Dict						= torch.manual_seed(0			)
        lowercase__    :   Optional[int]						= ddim(generator=__lowerCAmelCase     ,       eta=0.0     ,       output_type='''numpy'''			).images

        lowercase__    :   List[str]						= image[0, -3:, -3:, -1]

        assert image.shape == (1, 32, 32, 3)
        lowercase__    :   Optional[int]						= np.array([0.1_7_2_3, 0.1_6_1_7, 0.1_6_0_0, 0.1_6_2_6, 0.1_4_9_7, 0.1_5_1_3, 0.1_5_0_5, 0.1_4_4_2, 0.1_4_5_3]			)

        assert np.abs(image_slice.flatten() - expected_slice			).max() < 1E-2


       def _lowerCAmelCase(						self			)  ->  Union[str, Any]:
        lowercase__    :   Optional[Any]						= '''google/ddpm-ema-bedroom-256'''

        lowercase__    :   Tuple						= UNetaDModel.from_pretrained(__lowerCAmelCase			)
        lowercase__    :   Union[str, Any]						= DDIMScheduler.from_pretrained(__lowerCAmelCase			)

        lowercase__    :   List[Any]						= DDIMPipeline(unet=__lowerCAmelCase     ,       scheduler=__lowerCAmelCase			)
        ddpm.to(__lowerCAmelCase			)
        ddpm.set_progress_bar_config(disable=__lowerCAmelCase			)

        lowercase__    :   Union[str, Any]						= torch.manual_seed(0			)
        lowercase__    :   Optional[Any]						= ddpm(generator=__lowerCAmelCase     ,       output_type='''numpy'''			).images

        lowercase__    :   List[Any]						= image[0, -3:, -3:, -1]

        assert image.shape == (1, 256, 256, 3)
        lowercase__    :   Union[str, Any]						= np.array([0.0_0_6_0, 0.0_2_0_1, 0.0_3_4_4, 0.0_0_2_4, 0.0_0_1_8, 0.0_0_0_2, 0.0_0_2_2, 0.0_0_0_0, 0.0_0_6_9]			)

        assert np.abs(image_slice.flatten() - expected_slice			).max() < 1E-2

