




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

      def 	__magic_name__(   self		):
        lowerCAmelCase__				:  Optional[int]		       =     ids_tensor([self.batch_size, self.seq_length]      ,				self.vocab_size		)

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

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

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

        lowerCAmelCase__				:  Any		       =     self.get_config()

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

      def 	__magic_name__(   self		):
        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=__UpperCAmelCase      ,				initializer_range=self.initializer_range      ,				)

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

      def 	__magic_name__(   self      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				):
        lowerCAmelCase__				:  str		       =     True
        lowerCAmelCase__				:  int		       =     LlamaModel(__UpperCAmelCase		)
        model.to(__UpperCAmelCase		)
        model.eval()
        lowerCAmelCase__				:  Union[str, Any]		       =     model(
            __UpperCAmelCase      ,				attention_mask=__UpperCAmelCase      ,				encoder_hidden_states=__UpperCAmelCase      ,				encoder_attention_mask=__UpperCAmelCase      ,				)
        lowerCAmelCase__				:  int		       =     model(
            __UpperCAmelCase      ,				attention_mask=__UpperCAmelCase      ,				encoder_hidden_states=__UpperCAmelCase      ,				)
        lowerCAmelCase__				:  Optional[int]		       =     model(__UpperCAmelCase      ,				attention_mask=__UpperCAmelCase		)
        self.parent.assertEqual(result.last_hidden_state.shape      ,				(self.batch_size, self.seq_length, self.hidden_size)		)

      def 	__magic_name__(   self      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				):
        lowerCAmelCase__				:  str		       =     LlamaForCausalLM(config=__UpperCAmelCase		)
        model.to(__UpperCAmelCase		)
        model.eval()
        lowerCAmelCase__				:  Optional[Any]		       =     model(__UpperCAmelCase      ,				attention_mask=__UpperCAmelCase      ,				labels=__UpperCAmelCase		)
        self.parent.assertEqual(result.logits.shape      ,				(self.batch_size, self.seq_length, self.vocab_size)		)

      def 	__magic_name__(   self      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				):
        lowerCAmelCase__				:  int		       =     True
        lowerCAmelCase__				:  List[str]		       =     True
        lowerCAmelCase__				:  Optional[int]		       =     LlamaForCausalLM(config=__UpperCAmelCase		)
        model.to(__UpperCAmelCase		)
        model.eval()

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

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

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

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






@require_torch
class  _lowerCAmelCase     (  _lowercase    ,			_lowercase    ,			_lowercase    ,			unittest.TestCase	):
      A__		     =		(LlamaModel, LlamaForCausalLM, LlamaForSequenceClassification) if is_torch_available() else ()
      A__		     =		(LlamaForCausalLM,) if is_torch_available() else ()
      A__		     =		(
          {
              'feature-extraction': LlamaModel,
              'text-classification': LlamaForSequenceClassification,
              'text-generation': LlamaForCausalLM,
              'zero-shot': LlamaForSequenceClassification,
          }
          if is_torch_available()
          else {}
      )
      A__		     =		False
      A__		     =		False

      def 	__magic_name__(   self		):
        lowerCAmelCase__				:  Any		       =     LlamaModelTester(self		)
        lowerCAmelCase__				:  str		       =     ConfigTester(self      ,				config_class=__UpperCAmelCase      ,				hidden_size=37		)

      def 	__magic_name__(   self		):
        self.config_tester.run_common_tests()

      def 	__magic_name__(   self		):
        lowerCAmelCase__				:  int		       =     self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*__UpperCAmelCase		)

      def 	__magic_name__(   self		):
        lowerCAmelCase__				:  Union[str, Any]		       =     self.model_tester.prepare_config_and_inputs()
        for type in ["absolute", "relative_key", "relative_key_query"]:
          lowerCAmelCase__				:  int		       =     type
          self.model_tester.create_and_check_model(*__UpperCAmelCase		)

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

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

      def 	__magic_name__(   self		):
        lowerCAmelCase__       ,			lowerCAmelCase__				:  Union[str, Any]		       =     self.model_tester.prepare_config_and_inputs_for_common()
        lowerCAmelCase__				:  int		       =     3
        lowerCAmelCase__				:  Any		       =     '''multi_label_classification'''
        lowerCAmelCase__				:  int		       =     input_dict['''input_ids''']
        lowerCAmelCase__				:  Union[str, Any]		       =     input_ids.ne(1		).to(__UpperCAmelCase		)
        lowerCAmelCase__				:  Union[str, Any]		       =     ids_tensor(
            [self.model_tester.batch_size, config.num_labels]      ,				self.model_tester.type_sequence_label_size		).to(torch.float		)
        lowerCAmelCase__				:  Optional[Any]		       =     LlamaForSequenceClassification(__UpperCAmelCase		)
        model.to(__UpperCAmelCase		)
        model.eval()
        lowerCAmelCase__				:  Any		       =     model(__UpperCAmelCase      ,				attention_mask=__UpperCAmelCase      ,				labels=__UpperCAmelCase		)
        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 	__magic_name__(   self		):
        pass







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

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






@require_torch
class  _lowerCAmelCase     (  unittest.TestCase	):

      @unittest.skip('''Logits are not exactly the same, once we fix the instabalities somehow, will update!'''		)
      @slow
      def 	__magic_name__(   self		):
        lowerCAmelCase__				:  Tuple		       =     [1, 306, 4658, 278, 6593, 310, 2834, 338]
        lowerCAmelCase__				:  List[str]		       =     LlamaForCausalLM.from_pretrained('''meta-llama/Llama-2-7b-hf'''      ,				device_map='''auto'''		)
        lowerCAmelCase__				:  Any		       =     model(torch.tensor([input_ids]		)		)
        # Expected mean on dim = -1
        lowerCAmelCase__				:  Optional[int]		       =     torch.tensor([[-6.6550, -4.1227, -4.9859, -3.2406, 0.8262, -3.0033, 1.2964, -3.3699]]		)
        torch.testing.assert_close(out.mean(-1		)      ,				__UpperCAmelCase      ,				atol=1e-2      ,				rtol=1e-2		)
        # slicing logits[0, 0, 0:30]
        # fmt: off
        lowerCAmelCase__				:  List[Any]		       =     torch.tensor([-12.8281, -7.4453, -0.4639, -8.0625, -7.2500, -8.0000, -6.4883, -7.7695, -7.8438, -7.0312, -6.2188, -7.1328, -1.8496, 1.9961, -8.6250, -6.7227, -12.8281, -6.9492, -7.0742, -7.7852, -7.5820, -7.9062, -6.9375, -7.9805, -8.3438, -8.1562, -8.0469, -7.6250, -7.7422, -7.3398,]		)
        # fmt: on
        torch.testing.assert_close(out[0, 0, :30]      ,				__UpperCAmelCase      ,				atol=1e-5      ,				rtol=1e-5		)

      @unittest.skip('''Logits are not exactly the same, once we fix the instabalities somehow, will update!'''		)
      @slow
      def 	__magic_name__(   self		):
        lowerCAmelCase__				:  Optional[int]		       =     [1, 306, 4658, 278, 6593, 310, 2834, 338]
        lowerCAmelCase__				:  Dict		       =     LlamaForCausalLM.from_pretrained('''meta-llama/Llama-2-13b-hf'''      ,				device_map='''auto'''		)
        lowerCAmelCase__				:  Optional[int]		       =     model(torch.tensor(__UpperCAmelCase		)		)
        # Expected mean on dim = -1
        lowerCAmelCase__				:  Optional[int]		       =     torch.tensor([[-2.0622, -1.2794, -1.1638, -0.9788, -1.4603, -1.0238, -1.7893, -1.4411]]		)
        torch.testing.assert_close(out.mean(-1		)      ,				__UpperCAmelCase      ,				atol=1e-2      ,				rtol=1e-2		)
        # slicing logits[0, 0, 0:30]
        # fmt: off
        lowerCAmelCase__				:  Any		       =     torch.tensor([-8.1406, -8.0547, 2.7461, -1.2344, -0.1448, -1.8262, -1.0020, -1.8154, -1.6895, -1.8516, -2.3574, -0.9277, 3.7598, 6.5742, -1.2998, -0.1177, -8.1406, -2.9688, -2.9199, -3.1699, -3.5254, -2.3555, -2.7988, -3.4141, -2.8262, -4.5195, -3.3379, -3.3164, -2.7832, -3.0273]		)
        # fmt: on
        torch.testing.assert_close(out[0, 0, :30]      ,				__UpperCAmelCase      ,				atol=1e-5      ,				rtol=1e-5		)

      @unittest.skip('''Logits are not exactly the same, once we fix the instabalities somehow, will update!'''		)
      @slow
      def 	__magic_name__(   self		):
        lowerCAmelCase__				:  str		       =     [1, 306, 4658, 278, 6593, 310, 2834, 338]
        lowerCAmelCase__				:  Optional[Any]		       =     LlamaForCausalLM.from_pretrained('''meta-llama/Llama-2-13b-chat-hf'''      ,				device_map='''auto'''		)
        lowerCAmelCase__				:  List[str]		       =     model(torch.tensor(__UpperCAmelCase		)		)
        # Expected mean on dim = -1
        lowerCAmelCase__				:  Optional[Any]		       =     torch.tensor([[-0.8562, -1.8520, -0.7551, -0.4162, -1.5161, -1.2038, -2.4823, -2.3254]]		)
        torch.testing.assert_close(out.mean(-1		)      ,				__UpperCAmelCase      ,				atol=1e-2      ,				rtol=1e-2		)
        # slicing logits[0, 0, 0:30]
        # fmt: off
        lowerCAmelCase__				:  str		       =     torch.tensor([-2.2227, 4.8828, 0.9023, -0.4578, -0.7871, -0.1033, -0.6221, -0.5786, -0.7803, -1.0674, -1.2920, -0.1570, 0.8008, 2.0723, -0.9497, 0.2771, -2.2227, -0.7612, -1.4346, -1.2061, -1.6426, -0.3000, -0.7139, -1.1934, -1.8691, -1.6973, -1.5947, -1.2705, -0.3523, -0.5513]		)
        # fmt: on
        torch.testing.assert_close(out.mean(-1		)      ,				__UpperCAmelCase      ,				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 	__magic_name__(   self		):
        lowerCAmelCase__				:  Optional[Any]		       =     [1, 306, 4658, 278, 6593, 310, 2834, 338]
        lowerCAmelCase__				:  Optional[Any]		       =     LlamaForCausalLM.from_pretrained('''meta-llama/Llama-2-70b-hf'''      ,				device_map='''auto'''		)
        lowerCAmelCase__				:  Union[str, Any]		       =     model(torch.tensor(__UpperCAmelCase		)		)

        lowerCAmelCase__				:  Any		       =     torch.tensor(
            [[-4.2327, -3.3360, -4.6665, -4.7631, -1.8180, -3.4170, -1.4211, -3.1810]]      ,				dtype=torch.floataa		)
        torch.testing.assert_close(out.mean(-1		)      ,				__UpperCAmelCase      ,				atol=1e-2      ,				rtol=1e-2		)
        # fmt: off
        lowerCAmelCase__				:  int		       =     torch.tensor([-9.4922, -3.9551, 1.7998, -5.6758, -5.1055, -5.8984, -4.8320, -6.8086, -6.5391, -5.6172, -5.5820, -5.5352, 1.7881, 3.6289, -6.5117, -3.4785, -9.5000, -6.0352, -6.8125, -6.0195, -6.6836, -5.4727, -6.2812, -6.0391, -7.3398, -7.4297, -7.4844, -6.5820, -5.8789, -5.5312]		)
        # fmt: on
        torch.testing.assert_close(out[0, 0, :30]      ,				__UpperCAmelCase      ,				atol=1e-5      ,				rtol=1e-5		)







      @unittest.skip('''Model is curently gated'''		)
      @slow
      def 	__magic_name__(   self		):
        lowerCAmelCase__				:  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'''
        lowerCAmelCase__				:  int		       =     '''Simply put, the theory of relativity states that '''
        lowerCAmelCase__				:  int		       =     LlamaTokenizer.from_pretrained('''meta-llama/Llama-2-13b-chat-hf'''		)
        lowerCAmelCase__				:  Tuple		       =     tokenizer.encode(__UpperCAmelCase      ,				return_tensors='''pt'''		)
        lowerCAmelCase__				:  str		       =     LlamaForCausalLM.from_pretrained(
            '''meta-llama/Llama-2-13b-chat-hf'''      ,				device_map='''sequential'''      ,				use_safetensors=__UpperCAmelCase		)

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







from tempfile import TemporaryDirectory
from unittest import TestCase
from unittest.mock import MagicMock, patch

from transformers import AutoModel, TFAutoModel
from transformers.onnx import FeaturesManager
from transformers.testing_utils import SMALL_MODEL_IDENTIFIER, require_tf, require_torch






@require_torch
@require_tf
class  _lowerCAmelCase     (  _lowercase	):

      def 	__magic_name__(   self		):
        lowerCAmelCase__				:  Union[str, Any]		       =     SMALL_MODEL_IDENTIFIER
        lowerCAmelCase__				:  List[str]		       =     '''pt'''
        lowerCAmelCase__				:  int		       =     '''tf'''

      def 	__magic_name__(   self      ,				__UpperCAmelCase		):
        lowerCAmelCase__				:  Optional[Any]		       =     AutoModel.from_pretrained(self.test_model		)
        model_pt.save_pretrained(__UpperCAmelCase		)

      def 	__magic_name__(   self      ,				__UpperCAmelCase		):
        lowerCAmelCase__				:  str		       =     TFAutoModel.from_pretrained(self.test_model      ,				from_pt=__UpperCAmelCase		)
        model_tf.save_pretrained(__UpperCAmelCase		)

      def 	__magic_name__(   self		):
        lowerCAmelCase__				:  int		       =     '''mock_framework'''

        # Framework provided - return whatever the user provides
        lowerCAmelCase__				:  Tuple		       =     FeaturesManager.determine_framework(self.test_model      ,				__UpperCAmelCase		)
        self.assertEqual(__UpperCAmelCase      ,				__UpperCAmelCase		)

        # Local checkpoint and framework provided - return provided framework
        # PyTorch checkpoint
        with TemporaryDirectory() as local_pt_ckpt:
          self._setup_pt_ckpt(__UpperCAmelCase		)
          lowerCAmelCase__				:  int		       =     FeaturesManager.determine_framework(__UpperCAmelCase      ,				__UpperCAmelCase		)
          self.assertEqual(__UpperCAmelCase      ,				__UpperCAmelCase		)

        # TensorFlow checkpoint
        with TemporaryDirectory() as local_tf_ckpt:
          self._setup_tf_ckpt(__UpperCAmelCase		)
          lowerCAmelCase__				:  List[str]		       =     FeaturesManager.determine_framework(__UpperCAmelCase      ,				__UpperCAmelCase		)
          self.assertEqual(__UpperCAmelCase      ,				__UpperCAmelCase		)

      def 	__magic_name__(   self		):
        # PyTorch checkpoint
        with TemporaryDirectory() as local_pt_ckpt:
          self._setup_pt_ckpt(__UpperCAmelCase		)
          lowerCAmelCase__				:  Tuple		       =     FeaturesManager.determine_framework(__UpperCAmelCase		)
          self.assertEqual(__UpperCAmelCase      ,				self.framework_pt		)

        # TensorFlow checkpoint
        with TemporaryDirectory() as local_tf_ckpt:
          self._setup_tf_ckpt(__UpperCAmelCase		)
          lowerCAmelCase__				:  int		       =     FeaturesManager.determine_framework(__UpperCAmelCase		)
          self.assertEqual(__UpperCAmelCase      ,				self.framework_tf		)

        # Invalid local checkpoint
        with TemporaryDirectory() as local_invalid_ckpt:
          with self.assertRaises(__UpperCAmelCase		):
            lowerCAmelCase__				:  Tuple		       =     FeaturesManager.determine_framework(__UpperCAmelCase		)







      def 	__magic_name__(   self		):
        lowerCAmelCase__				:  List[Any]		       =     MagicMock(return_value=__UpperCAmelCase		)
        with patch('''transformers.onnx.features.is_tf_available'''      ,				__UpperCAmelCase		):
          lowerCAmelCase__				:  int		       =     FeaturesManager.determine_framework(self.test_model		)
          self.assertEqual(__UpperCAmelCase      ,				self.framework_pt		)

        # PyTorch not in environment -> use TensorFlow
        lowerCAmelCase__				:  Union[str, Any]		       =     MagicMock(return_value=__UpperCAmelCase		)
        with patch('''transformers.onnx.features.is_torch_available'''      ,				__UpperCAmelCase		):
          lowerCAmelCase__				:  Optional[int]		       =     FeaturesManager.determine_framework(self.test_model		)
          self.assertEqual(__UpperCAmelCase      ,				self.framework_tf		)

        # Both in environment -> use PyTorch
        lowerCAmelCase__				:  Union[str, Any]		       =     MagicMock(return_value=__UpperCAmelCase		)
        lowerCAmelCase__				:  List[Any]		       =     MagicMock(return_value=__UpperCAmelCase		)
        with patch('''transformers.onnx.features.is_tf_available'''      ,				__UpperCAmelCase		), patch(
            '''transformers.onnx.features.is_torch_available'''      ,				__UpperCAmelCase		):
          lowerCAmelCase__				:  List[Any]		       =     FeaturesManager.determine_framework(self.test_model		)
          self.assertEqual(__UpperCAmelCase      ,				self.framework_pt		)

        # Both not in environment -> raise error
        lowerCAmelCase__				:  Optional[int]		       =     MagicMock(return_value=__UpperCAmelCase		)
        lowerCAmelCase__				:  Dict		       =     MagicMock(return_value=__UpperCAmelCase		)
        with patch('''transformers.onnx.features.is_tf_available'''      ,				__UpperCAmelCase		), patch(
            '''transformers.onnx.features.is_torch_available'''      ,				__UpperCAmelCase		):
          with self.assertRaises(__UpperCAmelCase		):
            lowerCAmelCase__				:  Tuple		       =     FeaturesManager.determine_framework(self.test_model		)

