





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







      def __init__(    self		,      lowerCAmelCase__		,      lowerCAmelCase__=1_3		,      lowerCAmelCase__=7		,      lowerCAmelCase__=True		,      lowerCAmelCase__=True		,      lowerCAmelCase__=False		,      lowerCAmelCase__=True		,      lowerCAmelCase__=9_9		,      lowerCAmelCase__=3_2		,      lowerCAmelCase__=5		,      lowerCAmelCase__=4		,      lowerCAmelCase__=3_7		,      lowerCAmelCase__="gelu"		,      lowerCAmelCase__=0.1		,      lowerCAmelCase__=0.1		,      lowerCAmelCase__=5_1_2		,      lowerCAmelCase__=1_6		,      lowerCAmelCase__=2		,      lowerCAmelCase__=0.02		,      lowerCAmelCase__=3		,      lowerCAmelCase__=4		,      lowerCAmelCase__=None		,      )   ->						Optional[int]:



        '''simple docstring'''




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







      def    _lowercase  (    self							)   ->						Optional[int]:



        '''simple docstring'''




        a__	:					Dict										=ids_tensor([self.batch_size, self.seq_length]		,      self.vocab_size							)

        a__	:					Tuple										=None
        if self.use_input_mask:
          a__	:					Any										=random_attention_mask([self.batch_size, self.seq_length]							)

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

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

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

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







      def    _lowercase  (    self							)   ->						Optional[Any]:



        '''simple docstring'''




        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    _lowercase  (    self		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__							)   ->						List[Any]:



        '''simple docstring'''




        a__	:					str										=LlamaModel(config=lowerCAmelCase__							)
        model.to(lowerCAmelCase__							)
        model.eval()
        a__	:					Union[str, Any]										=model(lowerCAmelCase__		,      attention_mask=lowerCAmelCase__							)
        a__	:					Optional[Any]										=model(lowerCAmelCase__							)
        self.parent.assertEqual(result.last_hidden_state.shape		,      (self.batch_size, self.seq_length, self.hidden_size)							)







      def    _lowercase  (    self		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__		,      )   ->						Any:



        '''simple docstring'''




        a__	:					str										=True
        a__	:					Dict										=LlamaModel(lowerCAmelCase__							)
        model.to(lowerCAmelCase__							)
        model.eval()
        a__	:					List[Any]										=model(
            lowerCAmelCase__		,      attention_mask=lowerCAmelCase__		,      encoder_hidden_states=lowerCAmelCase__		,      encoder_attention_mask=lowerCAmelCase__		,      )
        a__	:					Optional[Any]										=model(
            lowerCAmelCase__		,      attention_mask=lowerCAmelCase__		,      encoder_hidden_states=lowerCAmelCase__		,      )
        a__	:					Dict										=model(lowerCAmelCase__		,      attention_mask=lowerCAmelCase__							)
        self.parent.assertEqual(result.last_hidden_state.shape		,      (self.batch_size, self.seq_length, self.hidden_size)							)







      def    _lowercase  (    self		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__		,      )   ->						Union[str, Any]:



        '''simple docstring'''




        a__	:					Dict										=LlamaForCausalLM(config=lowerCAmelCase__							)
        model.to(lowerCAmelCase__							)
        model.eval()
        a__	:					Dict										=model(lowerCAmelCase__		,      attention_mask=lowerCAmelCase__		,      labels=lowerCAmelCase__							)
        self.parent.assertEqual(result.logits.shape		,      (self.batch_size, self.seq_length, self.vocab_size)							)







      def    _lowercase  (    self		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__		,      )   ->						List[Any]:



        '''simple docstring'''




        a__	:					List[Any]										=True
        a__	:					Optional[Any]										=True
        a__	:					List[Any]										=LlamaForCausalLM(config=lowerCAmelCase__							)
        model.to(lowerCAmelCase__							)
        model.eval()

        # first forward pass
        a__	:					int										=model(
            lowerCAmelCase__		,      attention_mask=lowerCAmelCase__		,      encoder_hidden_states=lowerCAmelCase__		,      encoder_attention_mask=lowerCAmelCase__		,      use_cache=lowerCAmelCase__		,      )
        a__	:					List[Any]										=outputs.past_key_values

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

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

        a__	:					int										=model(
            lowerCAmelCase__		,      attention_mask=lowerCAmelCase__		,      encoder_hidden_states=lowerCAmelCase__		,      encoder_attention_mask=lowerCAmelCase__		,      output_hidden_states=lowerCAmelCase__		,      )["hidden_states"][0]
        a__	:					Union[str, Any]										=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
        a__	:					Optional[Any]										=ids_tensor((1,)		,      output_from_past.shape[-1]							).item()
        a__	:					List[str]										=output_from_no_past[:, -3:, random_slice_idx].detach()
        a__	:					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(lowerCAmelCase__		,      lowerCAmelCase__		,      atol=1E-3							)							)







      def    _lowercase  (    self							)   ->						Any:



        '''simple docstring'''




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







@require_torch
class 				__lowerCAmelCase       (    UpperCamelCase__						,				UpperCamelCase__						,				UpperCamelCase__						,				unittest.TestCase):
      _lowercase				:					Dict						=			(LlamaModel, LlamaForCausalLM, LlamaForSequenceClassification) if is_torch_available() else ()
      _lowercase				:					Union[str, Any]						=			(LlamaForCausalLM,) if is_torch_available() else ()
      _lowercase				:					List[Any]						=			(
          {
              """feature-extraction""": LlamaModel,
              """text-classification""": LlamaForSequenceClassification,
              """text-generation""": LlamaForCausalLM,
              """zero-shot""": LlamaForSequenceClassification,
          }
          if is_torch_available()
          else {}
      )
      _lowercase				:					List[Any]						=			False
      _lowercase				:					Optional[Any]						=			False







      def    _lowercase  (    self							)   ->						Optional[int]:



        '''simple docstring'''




        a__	:					Optional[int]										=LlamaModelTester(self							)
        a__	:					Dict										=ConfigTester(self		,      config_class=lowerCAmelCase__		,      hidden_size=3_7							)







      def    _lowercase  (    self							)   ->						Dict:



        '''simple docstring'''




        self.config_tester.run_common_tests()







      def    _lowercase  (    self							)   ->						Union[str, Any]:



        '''simple docstring'''




        a__	:					List[str]										=self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*lowerCAmelCase__							)







      def    _lowercase  (    self							)   ->						List[Any]:



        '''simple docstring'''




        a__	:					Dict										=self.model_tester.prepare_config_and_inputs()
        for type in ["absolute", "relative_key", "relative_key_query"]:
          a__	:					Dict										=type
          self.model_tester.create_and_check_model(*lowerCAmelCase__							)







      def    _lowercase  (    self							)   ->						Union[str, Any]:



        '''simple docstring'''




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







      def    _lowercase  (    self							)   ->						str:



        '''simple docstring'''




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







      def    _lowercase  (    self							)   ->						Dict:



        '''simple docstring'''




        a__					,							a__	:					Optional[Any]										=self.model_tester.prepare_config_and_inputs_for_common()
        a__	:					Tuple										=3
        a__	:					List[str]										="multi_label_classification"
        a__	:					Optional[int]										=input_dict["input_ids"]
        a__	:					Dict										=input_ids.ne(1							).to(lowerCAmelCase__							)
        a__	:					List[Any]										=ids_tensor(
            [self.model_tester.batch_size, config.num_labels]		,      self.model_tester.type_sequence_label_size							).to(torch.float							)
        a__	:					List[str]										=LlamaForSequenceClassification(lowerCAmelCase__							)
        model.to(lowerCAmelCase__							)
        model.eval()
        a__	:					int										=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    _lowercase  (    self							)   ->						Dict:



        '''simple docstring'''




        pass







      @parameterized.expand([("linear",), ("dynamic",)]							)
      def    _lowercase  (    self		,      lowerCAmelCase__							)   ->						Optional[Any]:



        '''simple docstring'''




        a__					,							a__	:					Optional[int]										=self.model_tester.prepare_config_and_inputs_for_common()
        a__	:					Tuple										=ids_tensor([1, 1_0]		,      config.vocab_size							)
        a__	:					Optional[int]										=ids_tensor([1, int(config.max_position_embeddings * 1.5							)]		,      config.vocab_size							)

        set_seed(4_2							)  # Fixed seed at init time so the two models get the same random weights
        a__	:					List[Any]										=LlamaModel(lowerCAmelCase__							)
        original_model.to(lowerCAmelCase__							)
        original_model.eval()
        a__	:					Union[str, Any]										=original_model(lowerCAmelCase__							).last_hidden_state
        a__	:					Dict										=original_model(lowerCAmelCase__							).last_hidden_state

        set_seed(4_2							)  # Fixed seed at init time so the two models get the same random weights
        a__	:					List[str]										={"type": scaling_type, "factor": 10.0}
        a__	:					Optional[Any]										=LlamaModel(lowerCAmelCase__							)
        scaled_model.to(lowerCAmelCase__							)
        scaled_model.eval()
        a__	:					List[str]										=scaled_model(lowerCAmelCase__							).last_hidden_state
        a__	:					List[str]										=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 				__lowerCAmelCase       (    unittest.TestCase):







      @unittest.skip("Logits are not exactly the same, once we fix the instabalities somehow, will update!"							)
      @slow
      def    _lowercase  (    self							)   ->						Tuple:



        '''simple docstring'''




        a__	:					int										=[1, 3_0_6, 4_6_5_8, 2_7_8, 6_5_9_3, 3_1_0, 2_8_3_4, 3_3_8]
        a__	:					Dict										=LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf"		,      device_map="auto"							)
        a__	:					Optional[Any]										=model(torch.tensor([input_ids]							)							)
        # Expected mean on dim = -1
        a__	:					Any										=torch.tensor([[-6.65_50, -4.12_27, -4.98_59, -3.24_06, 0.82_62, -3.00_33, 1.29_64, -3.36_99]]							)
        torch.testing.assert_close(out.mean(-1							)		,      lowerCAmelCase__		,      atol=1E-2		,      rtol=1E-2							)
        # slicing logits[0, 0, 0:30]
        # fmt: off
        a__	:					int										=torch.tensor([-12.82_81, -7.44_53, -0.46_39, -8.06_25, -7.25_00, -8.00_00, -6.48_83, -7.76_95, -7.84_38, -7.03_12, -6.21_88, -7.13_28, -1.84_96, 1.99_61, -8.62_50, -6.72_27, -12.82_81, -6.94_92, -7.07_42, -7.78_52, -7.58_20, -7.90_62, -6.93_75, -7.98_05, -8.34_38, -8.15_62, -8.04_69, -7.62_50, -7.74_22, -7.33_98,]							)
        # fmt: on
        torch.testing.assert_close(out[0, 0, :3_0]		,      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    _lowercase  (    self							)   ->						List[str]:



        '''simple docstring'''




        a__	:					Tuple										=[1, 3_0_6, 4_6_5_8, 2_7_8, 6_5_9_3, 3_1_0, 2_8_3_4, 3_3_8]
        a__	:					int										=LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-13b-hf"		,      device_map="auto"							)
        a__	:					List[Any]										=model(torch.tensor(lowerCAmelCase__							)							)
        # Expected mean on dim = -1
        a__	:					Any										=torch.tensor([[-2.06_22, -1.27_94, -1.16_38, -0.97_88, -1.46_03, -1.02_38, -1.78_93, -1.44_11]]							)
        torch.testing.assert_close(out.mean(-1							)		,      lowerCAmelCase__		,      atol=1E-2		,      rtol=1E-2							)
        # slicing logits[0, 0, 0:30]
        # fmt: off
        a__	:					Optional[int]										=torch.tensor([-8.14_06, -8.05_47, 2.74_61, -1.23_44, -0.14_48, -1.82_62, -1.00_20, -1.81_54, -1.68_95, -1.85_16, -2.35_74, -0.92_77, 3.75_98, 6.57_42, -1.29_98, -0.11_77, -8.14_06, -2.96_88, -2.91_99, -3.16_99, -3.52_54, -2.35_55, -2.79_88, -3.41_41, -2.82_62, -4.51_95, -3.33_79, -3.31_64, -2.78_32, -3.02_73]							)
        # fmt: on
        torch.testing.assert_close(out[0, 0, :3_0]		,      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    _lowercase  (    self							)   ->						int:



        '''simple docstring'''




        a__	:					int										=[1, 3_0_6, 4_6_5_8, 2_7_8, 6_5_9_3, 3_1_0, 2_8_3_4, 3_3_8]
        a__	:					Optional[int]										=LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-13b-chat-hf"		,      device_map="auto"							)
        a__	:					Union[str, Any]										=model(torch.tensor(lowerCAmelCase__							)							)
        # Expected mean on dim = -1
        a__	:					Tuple										=torch.tensor([[-0.85_62, -1.85_20, -0.75_51, -0.41_62, -1.51_61, -1.20_38, -2.48_23, -2.32_54]]							)
        torch.testing.assert_close(out.mean(-1							)		,      lowerCAmelCase__		,      atol=1E-2		,      rtol=1E-2							)
        # slicing logits[0, 0, 0:30]
        # fmt: off
        a__	:					List[str]										=torch.tensor([-2.22_27, 4.88_28, 0.90_23, -0.45_78, -0.78_71, -0.10_33, -0.62_21, -0.57_86, -0.78_03, -1.06_74, -1.29_20, -0.15_70, 0.80_08, 2.07_23, -0.94_97, 0.27_71, -2.22_27, -0.76_12, -1.43_46, -1.20_61, -1.64_26, -0.30_00, -0.71_39, -1.19_34, -1.86_91, -1.69_73, -1.59_47, -1.27_05, -0.35_23, -0.55_13]							)
        # 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    _lowercase  (    self							)   ->						Dict:



        '''simple docstring'''




        a__	:					Tuple										=[1, 3_0_6, 4_6_5_8, 2_7_8, 6_5_9_3, 3_1_0, 2_8_3_4, 3_3_8]
        a__	:					List[Any]										=LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-70b-hf"		,      device_map="auto"							)
        a__	:					List[Any]										=model(torch.tensor(lowerCAmelCase__							)							)

        a__	:					Tuple										=torch.tensor(
            [[-4.23_27, -3.33_60, -4.66_65, -4.76_31, -1.81_80, -3.41_70, -1.42_11, -3.18_10]]		,      dtype=torch.floataa							)
        torch.testing.assert_close(out.mean(-1							)		,      lowerCAmelCase__		,      atol=1E-2		,      rtol=1E-2							)
        # fmt: off
        a__	:					Union[str, Any]										=torch.tensor([-9.49_22, -3.95_51, 1.79_98, -5.67_58, -5.10_55, -5.89_84, -4.83_20, -6.80_86, -6.53_91, -5.61_72, -5.58_20, -5.53_52, 1.78_81, 3.62_89, -6.51_17, -3.47_85, -9.50_00, -6.03_52, -6.81_25, -6.01_95, -6.68_36, -5.47_27, -6.28_12, -6.03_91, -7.33_98, -7.42_97, -7.48_44, -6.58_20, -5.87_89, -5.53_12]							)
        # fmt: on
        torch.testing.assert_close(out[0, 0, :3_0]		,      lowerCAmelCase__		,      atol=1E-5		,      rtol=1E-5							)







      @unittest.skip("Model is curently gated"							)
      @slow
      def    _lowercase  (    self							)   ->						int:



        '''simple docstring'''




        a__	:					List[Any]										="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"
        a__	:					str										="Simply put, the theory of relativity states that "
        a__	:					Tuple										=LlamaTokenizer.from_pretrained("meta-llama/Llama-2-13b-chat-hf"							)
        a__	:					Optional[Any]										=tokenizer.encode(lowerCAmelCase__		,      return_tensors="pt"							)
        a__	:					Union[str, Any]										=LlamaForCausalLM.from_pretrained(
            "meta-llama/Llama-2-13b-chat-hf"		,      device_map="sequential"		,      use_safetensors=lowerCAmelCase__							)

        # greedy generation outputs
        a__	:					Union[str, Any]										=model.generate(lowerCAmelCase__		,      max_new_tokens=6_4		,      top_p=lowerCAmelCase__		,      temperature=1		,      do_sample=lowerCAmelCase__							)
        a__	:					int										=tokenizer.decode(generated_ids[0]		,      skip_special_tokens=lowerCAmelCase__							)
        self.assertEqual(lowerCAmelCase__		,      lowerCAmelCase__							)










# DISCLAIMER: This file is strongly influenced by https://github.com/yang-song/score_sde_pytorch

import math
from typing import Union

import torch

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import randn_tensor
from .scheduling_utils import SchedulerMixin







class       lowerCamelCase__   (       lowerCAmelCase   ,							lowerCAmelCase):
     SCREAMING_SNAKE_CASE__      			=			1





     @register_to_config
     def __init__(self		,     UpperCAmelCase=2_0_0_0		,     UpperCAmelCase=0.1		,     UpperCAmelCase=2_0		,     UpperCAmelCase=1e-3       )     ->   List[str]:
      _lowercase							      =None
      _lowercase							      =None
      _lowercase							      =None





     def      __A					(self		,     UpperCAmelCase		,     UpperCAmelCase = None       )     ->   str:
      _lowercase							      =torch.linspace(1		,     self.config.sampling_eps		,     UpperCAmelCase		,     device=UpperCAmelCase       )





     def      __A					(self		,     UpperCAmelCase		,     UpperCAmelCase		,     UpperCAmelCase		,     UpperCAmelCase=None       )     ->   Optional[int]:
      if self.timesteps is None:
       raise ValueError(
           '''`self.timesteps` is not set, you need to run \'set_timesteps\' after creating the scheduler'''       )

      # TODO(Patrick) better comments + non-PyTorch
      # postprocess model score
      _lowercase							      =(
          -0.25 * t**2 * (self.config.beta_max - self.config.beta_min) - 0.5 * t * self.config.beta_min
      )
      _lowercase							      =torch.sqrt(1.0 - torch.exp(2.0 * log_mean_coeff       )       )
      _lowercase							      =std.flatten()
      while len(std.shape       ) < len(score.shape       ):
       _lowercase							      =std.unsqueeze(-1       )
      _lowercase							      =-score / std

      # compute
      _lowercase							      =-1.0 / len(self.timesteps       )

      _lowercase							      =self.config.beta_min + t * (self.config.beta_max - self.config.beta_min)
      _lowercase							      =beta_t.flatten()
      while len(beta_t.shape       ) < len(x.shape       ):
       _lowercase							      =beta_t.unsqueeze(-1       )
      _lowercase							      =-0.5 * beta_t * x

      _lowercase							      =torch.sqrt(UpperCAmelCase       )
      _lowercase							      =drift - diffusion**2 * score
      _lowercase							      =x + drift * dt

      # add noise
      _lowercase							      =randn_tensor(x.shape		,     layout=x.layout		,     generator=UpperCAmelCase		,     device=x.device		,     dtype=x.dtype       )
      _lowercase							      =x_mean + diffusion * math.sqrt(-dt       ) * noise

      return x, x_mean





     def __len__(self       )     ->   str:
      return self.config.num_train_timesteps

