

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

     SCREAMING_SNAKE_CASE:Optional[int]						      =					None
     if self.use_input_mask:
       SCREAMING_SNAKE_CASE:str						      =					random_attention_mask([self.batch_size, self.seq_length]							)

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

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

     SCREAMING_SNAKE_CASE:Optional[int]						      =					self.get_config()

     return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
   def     __UpperCamelCase						( self			:      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=SCREAMING_SNAKE_CASE__     ,initializer_range=self.initializer_range     ,)
   def     __UpperCamelCase						( self			:      Tuple     ,SCREAMING_SNAKE_CASE__			:      Tuple     ,SCREAMING_SNAKE_CASE__			:      Tuple     ,SCREAMING_SNAKE_CASE__			:      int     ,SCREAMING_SNAKE_CASE__			:      List[str]     ,SCREAMING_SNAKE_CASE__			:      List[str]     ,SCREAMING_SNAKE_CASE__			:      Optional[Any]     ,SCREAMING_SNAKE_CASE__			:      Optional[Any]							):
     SCREAMING_SNAKE_CASE:Union[str, Any]						      =					LlamaModel(config=SCREAMING_SNAKE_CASE__							)
     model.to(SCREAMING_SNAKE_CASE__							)
     model.eval()
     SCREAMING_SNAKE_CASE:Dict						      =					model(SCREAMING_SNAKE_CASE__     ,attention_mask=SCREAMING_SNAKE_CASE__							)
     SCREAMING_SNAKE_CASE:int						      =					model(SCREAMING_SNAKE_CASE__							)
     self.parent.assertEqual(result.last_hidden_state.shape     ,(self.batch_size, self.seq_length, self.hidden_size)							)
   def     __UpperCamelCase						( self			:      str     ,SCREAMING_SNAKE_CASE__			:      Dict     ,SCREAMING_SNAKE_CASE__			:      List[str]     ,SCREAMING_SNAKE_CASE__			:      Union[str, Any]     ,SCREAMING_SNAKE_CASE__			:      List[str]     ,SCREAMING_SNAKE_CASE__			:      Dict     ,SCREAMING_SNAKE_CASE__			:      Union[str, Any]     ,SCREAMING_SNAKE_CASE__			:      Dict     ,SCREAMING_SNAKE_CASE__			:      Union[str, Any]     ,SCREAMING_SNAKE_CASE__			:      Dict     ,):
     SCREAMING_SNAKE_CASE:Tuple						      =					True
     SCREAMING_SNAKE_CASE:Union[str, Any]						      =					LlamaModel(SCREAMING_SNAKE_CASE__							)
     model.to(SCREAMING_SNAKE_CASE__							)
     model.eval()
     SCREAMING_SNAKE_CASE:Tuple						      =					model(
         SCREAMING_SNAKE_CASE__     ,attention_mask=SCREAMING_SNAKE_CASE__     ,encoder_hidden_states=SCREAMING_SNAKE_CASE__     ,encoder_attention_mask=SCREAMING_SNAKE_CASE__     ,)
     SCREAMING_SNAKE_CASE:Optional[Any]						      =					model(
         SCREAMING_SNAKE_CASE__     ,attention_mask=SCREAMING_SNAKE_CASE__     ,encoder_hidden_states=SCREAMING_SNAKE_CASE__     ,)
     SCREAMING_SNAKE_CASE:Union[str, Any]						      =					model(SCREAMING_SNAKE_CASE__     ,attention_mask=SCREAMING_SNAKE_CASE__							)
     self.parent.assertEqual(result.last_hidden_state.shape     ,(self.batch_size, self.seq_length, self.hidden_size)							)
   def     __UpperCamelCase						( self			:      Union[str, Any]     ,SCREAMING_SNAKE_CASE__			:      Any     ,SCREAMING_SNAKE_CASE__			:      Tuple     ,SCREAMING_SNAKE_CASE__			:      Optional[Any]     ,SCREAMING_SNAKE_CASE__			:      Optional[Any]     ,SCREAMING_SNAKE_CASE__			:      List[str]     ,SCREAMING_SNAKE_CASE__			:      Any     ,SCREAMING_SNAKE_CASE__			:      Dict     ,SCREAMING_SNAKE_CASE__			:      Union[str, Any]     ,SCREAMING_SNAKE_CASE__			:      Optional[int]     ,):
     SCREAMING_SNAKE_CASE:Union[str, Any]						      =					LlamaForCausalLM(config=SCREAMING_SNAKE_CASE__							)
     model.to(SCREAMING_SNAKE_CASE__							)
     model.eval()
     SCREAMING_SNAKE_CASE:Union[str, Any]						      =					model(SCREAMING_SNAKE_CASE__     ,attention_mask=SCREAMING_SNAKE_CASE__     ,labels=SCREAMING_SNAKE_CASE__							)
     self.parent.assertEqual(result.logits.shape     ,(self.batch_size, self.seq_length, self.vocab_size)							)
   def     __UpperCamelCase						( self			:      int     ,SCREAMING_SNAKE_CASE__			:      List[Any]     ,SCREAMING_SNAKE_CASE__			:      Any     ,SCREAMING_SNAKE_CASE__			:      List[str]     ,SCREAMING_SNAKE_CASE__			:      str     ,SCREAMING_SNAKE_CASE__			:      Union[str, Any]     ,SCREAMING_SNAKE_CASE__			:      str     ,SCREAMING_SNAKE_CASE__			:      str     ,SCREAMING_SNAKE_CASE__			:      str     ,SCREAMING_SNAKE_CASE__			:      Any     ,):
     SCREAMING_SNAKE_CASE:Any						      =					True
     SCREAMING_SNAKE_CASE:Dict						      =					True
     SCREAMING_SNAKE_CASE:Tuple						      =					LlamaForCausalLM(config=SCREAMING_SNAKE_CASE__							)
     model.to(SCREAMING_SNAKE_CASE__							)
     model.eval()

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

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

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

     SCREAMING_SNAKE_CASE:Optional[int]						      =					model(
         SCREAMING_SNAKE_CASE__     ,attention_mask=SCREAMING_SNAKE_CASE__     ,encoder_hidden_states=SCREAMING_SNAKE_CASE__     ,encoder_attention_mask=SCREAMING_SNAKE_CASE__     ,output_hidden_states=SCREAMING_SNAKE_CASE__     ,)["hidden_states"][0]
     SCREAMING_SNAKE_CASE:Tuple						      =					model(
         SCREAMING_SNAKE_CASE__     ,attention_mask=SCREAMING_SNAKE_CASE__     ,encoder_hidden_states=SCREAMING_SNAKE_CASE__     ,encoder_attention_mask=SCREAMING_SNAKE_CASE__     ,past_key_values=SCREAMING_SNAKE_CASE__     ,output_hidden_states=SCREAMING_SNAKE_CASE__     ,)["hidden_states"][0]

     # select random slice
     SCREAMING_SNAKE_CASE:Optional[Any]						      =					ids_tensor((1,)     ,output_from_past.shape[-1]							).item()
     SCREAMING_SNAKE_CASE:List[Any]						      =					output_from_no_past[:, -3:, random_slice_idx].detach()
     SCREAMING_SNAKE_CASE: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(SCREAMING_SNAKE_CASE__     ,SCREAMING_SNAKE_CASE__     ,atol=1e-3							)							)







   def     __UpperCamelCase						( self			:      str							):
     SCREAMING_SNAKE_CASE:List[str]						      =					self.prepare_config_and_inputs()
     (
         (
         SCREAMING_SNAKE_CASE
     )      ,    (
         SCREAMING_SNAKE_CASE
     )      ,    (
         SCREAMING_SNAKE_CASE
     )      ,    (
         SCREAMING_SNAKE_CASE
     )      ,    (
         SCREAMING_SNAKE_CASE
     )      ,    (
         SCREAMING_SNAKE_CASE
     )      ,    (
         SCREAMING_SNAKE_CASE
     )      ,    
     ):int						      =					config_and_inputs
     SCREAMING_SNAKE_CASE:Tuple						      =					{"input_ids": input_ids, "attention_mask": input_mask}
     return config, inputs_dict





@require_torch
class 							_snake_case (       _a	,  _a	,  _a	,  unittest.TestCase						):
   _A       :     Union[str, Any]       =					(LlamaModel, LlamaForCausalLM, LlamaForSequenceClassification) if is_torch_available() else ()
   _A       :     Dict       =					(LlamaForCausalLM,) if is_torch_available() else ()
   _A       :     Union[str, Any]       =					(
       {
           '''feature-extraction''': LlamaModel,
           '''text-classification''': LlamaForSequenceClassification,
           '''text-generation''': LlamaForCausalLM,
           '''zero-shot''': LlamaForSequenceClassification,
       }
       if is_torch_available()
       else {}
   )
   _A       :     Tuple       =					False
   _A       :     Any       =					False
   def     __UpperCamelCase						( self			:      int							):
     SCREAMING_SNAKE_CASE:Optional[int]						      =					LlamaModelTester(self							)
     SCREAMING_SNAKE_CASE:int						      =					ConfigTester(self     ,config_class=SCREAMING_SNAKE_CASE__     ,hidden_size=37							)
   def     __UpperCamelCase						( self			:      str							):
     self.config_tester.run_common_tests()
   def     __UpperCamelCase						( self			:      Optional[Any]							):
     SCREAMING_SNAKE_CASE:List[str]						      =					self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE__							)
   def     __UpperCamelCase						( self			:      Dict							):
     SCREAMING_SNAKE_CASE:str						      =					self.model_tester.prepare_config_and_inputs()
     for type in ["absolute", "relative_key", "relative_key_query"]:
       SCREAMING_SNAKE_CASE:Dict						      =					type
       self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE__							)
   def     __UpperCamelCase						( self			:      Union[str, Any]							):
     SCREAMING_SNAKE_CASE      ,    SCREAMING_SNAKE_CASE:Union[str, Any]						      =					self.model_tester.prepare_config_and_inputs_for_common()
     SCREAMING_SNAKE_CASE:Optional[int]						      =					3
     SCREAMING_SNAKE_CASE:Tuple						      =					input_dict["input_ids"]
     SCREAMING_SNAKE_CASE:Optional[int]						      =					input_ids.ne(1							).to(SCREAMING_SNAKE_CASE__							)
     SCREAMING_SNAKE_CASE:Any						      =					ids_tensor([self.model_tester.batch_size]     ,self.model_tester.type_sequence_label_size							)
     SCREAMING_SNAKE_CASE:Dict						      =					LlamaForSequenceClassification(SCREAMING_SNAKE_CASE__							)
     model.to(SCREAMING_SNAKE_CASE__							)
     model.eval()
     SCREAMING_SNAKE_CASE:Dict						      =					model(SCREAMING_SNAKE_CASE__     ,attention_mask=SCREAMING_SNAKE_CASE__     ,labels=SCREAMING_SNAKE_CASE__							)
     self.assertEqual(result.logits.shape     ,(self.model_tester.batch_size, self.model_tester.num_labels)							)
   def     __UpperCamelCase						( self			:      str							):
     SCREAMING_SNAKE_CASE      ,    SCREAMING_SNAKE_CASE:Tuple						      =					self.model_tester.prepare_config_and_inputs_for_common()
     SCREAMING_SNAKE_CASE:int						      =					3
     SCREAMING_SNAKE_CASE:Any						      =					"single_label_classification"
     SCREAMING_SNAKE_CASE:Union[str, Any]						      =					input_dict["input_ids"]
     SCREAMING_SNAKE_CASE:Optional[int]						      =					input_ids.ne(1							).to(SCREAMING_SNAKE_CASE__							)
     SCREAMING_SNAKE_CASE:List[str]						      =					ids_tensor([self.model_tester.batch_size]     ,self.model_tester.type_sequence_label_size							)
     SCREAMING_SNAKE_CASE:Optional[Any]						      =					LlamaForSequenceClassification(SCREAMING_SNAKE_CASE__							)
     model.to(SCREAMING_SNAKE_CASE__							)
     model.eval()
     SCREAMING_SNAKE_CASE:Dict						      =					model(SCREAMING_SNAKE_CASE__     ,attention_mask=SCREAMING_SNAKE_CASE__     ,labels=SCREAMING_SNAKE_CASE__							)
     self.assertEqual(result.logits.shape     ,(self.model_tester.batch_size, self.model_tester.num_labels)							)
   def     __UpperCamelCase						( self			:      Tuple							):
     SCREAMING_SNAKE_CASE      ,    SCREAMING_SNAKE_CASE:Optional[int]						      =					self.model_tester.prepare_config_and_inputs_for_common()
     SCREAMING_SNAKE_CASE:Tuple						      =					3
     SCREAMING_SNAKE_CASE:Any						      =					"multi_label_classification"
     SCREAMING_SNAKE_CASE:str						      =					input_dict["input_ids"]
     SCREAMING_SNAKE_CASE:str						      =					input_ids.ne(1							).to(SCREAMING_SNAKE_CASE__							)
     SCREAMING_SNAKE_CASE:Dict						      =					ids_tensor(
         [self.model_tester.batch_size, config.num_labels]     ,self.model_tester.type_sequence_label_size							).to(torch.float							)
     SCREAMING_SNAKE_CASE:Dict						      =					LlamaForSequenceClassification(SCREAMING_SNAKE_CASE__							)
     model.to(SCREAMING_SNAKE_CASE__							)
     model.eval()
     SCREAMING_SNAKE_CASE:Dict						      =					model(SCREAMING_SNAKE_CASE__     ,attention_mask=SCREAMING_SNAKE_CASE__     ,labels=SCREAMING_SNAKE_CASE__							)
     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     __UpperCamelCase						( self			:      Dict							):
     pass







   @parameterized.expand([("linear",), ("dynamic",)]							)
   def     __UpperCamelCase						( self			:      Dict     ,SCREAMING_SNAKE_CASE__			:      Tuple							):
     SCREAMING_SNAKE_CASE      ,    SCREAMING_SNAKE_CASE:List[str]						      =					self.model_tester.prepare_config_and_inputs_for_common()
     SCREAMING_SNAKE_CASE:Any						      =					ids_tensor([1, 10]     ,config.vocab_size							)
     SCREAMING_SNAKE_CASE:List[str]						      =					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
     SCREAMING_SNAKE_CASE:List[str]						      =					LlamaModel(SCREAMING_SNAKE_CASE__							)
     original_model.to(SCREAMING_SNAKE_CASE__							)
     original_model.eval()
     SCREAMING_SNAKE_CASE:Dict						      =					original_model(SCREAMING_SNAKE_CASE__							).last_hidden_state
     SCREAMING_SNAKE_CASE:List[Any]						      =					original_model(SCREAMING_SNAKE_CASE__							).last_hidden_state

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

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





@require_torch
class 							_snake_case (       unittest.TestCase						):
   @unittest.skip("Logits are not exactly the same, once we fix the instabalities somehow, will update!"							)
   @slow
   def     __UpperCamelCase						( self			:      List[str]							):
     SCREAMING_SNAKE_CASE:Dict						      =					[1, 306, 4_658, 278, 6_593, 310, 2_834, 338]
     SCREAMING_SNAKE_CASE:List[str]						      =					LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf"     ,device_map="auto"							)
     SCREAMING_SNAKE_CASE:Tuple						      =					model(torch.tensor([input_ids]							)							)
     # Expected mean on dim = -1
     SCREAMING_SNAKE_CASE:str						      =					torch.tensor([[-6.6_550, -4.1_227, -4.9_859, -3.2_406, 0.8_262, -3.0_033, 1.2_964, -3.3_699]]							)
     torch.testing.assert_close(out.mean(-1							)     ,SCREAMING_SNAKE_CASE__     ,atol=1e-2     ,rtol=1e-2							)
     # slicing logits[0, 0, 0:30]
     # fmt: off
     SCREAMING_SNAKE_CASE:Optional[Any]						      =					torch.tensor([-12.8_281, -7.4_453, -0.4_639, -8.0_625, -7.2_500, -8.0_000, -6.4_883, -7.7_695, -7.8_438, -7.0_312, -6.2_188, -7.1_328, -1.8_496, 1.9_961, -8.6_250, -6.7_227, -12.8_281, -6.9_492, -7.0_742, -7.7_852, -7.5_820, -7.9_062, -6.9_375, -7.9_805, -8.3_438, -8.1_562, -8.0_469, -7.6_250, -7.7_422, -7.3_398,]							)
     # fmt: on
     torch.testing.assert_close(out[0, 0, :30]     ,SCREAMING_SNAKE_CASE__     ,atol=1e-5     ,rtol=1e-5							)
   @unittest.skip("Logits are not exactly the same, once we fix the instabalities somehow, will update!"							)
   @slow
   def     __UpperCamelCase						( self			:      List[Any]							):
     SCREAMING_SNAKE_CASE:Tuple						      =					[1, 306, 4_658, 278, 6_593, 310, 2_834, 338]
     SCREAMING_SNAKE_CASE:Tuple						      =					LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-13b-hf"     ,device_map="auto"							)
     SCREAMING_SNAKE_CASE:Dict						      =					model(torch.tensor(SCREAMING_SNAKE_CASE__							)							)
     # Expected mean on dim = -1
     SCREAMING_SNAKE_CASE:Dict						      =					torch.tensor([[-2.0_622, -1.2_794, -1.1_638, -0.9_788, -1.4_603, -1.0_238, -1.7_893, -1.4_411]]							)
     torch.testing.assert_close(out.mean(-1							)     ,SCREAMING_SNAKE_CASE__     ,atol=1e-2     ,rtol=1e-2							)
     # slicing logits[0, 0, 0:30]
     # fmt: off
     SCREAMING_SNAKE_CASE:Tuple						      =					torch.tensor([-8.1_406, -8.0_547, 2.7_461, -1.2_344, -0.1_448, -1.8_262, -1.0_020, -1.8_154, -1.6_895, -1.8_516, -2.3_574, -0.9_277, 3.7_598, 6.5_742, -1.2_998, -0.1_177, -8.1_406, -2.9_688, -2.9_199, -3.1_699, -3.5_254, -2.3_555, -2.7_988, -3.4_141, -2.8_262, -4.5_195, -3.3_379, -3.3_164, -2.7_832, -3.0_273]							)
     # fmt: on
     torch.testing.assert_close(out[0, 0, :30]     ,SCREAMING_SNAKE_CASE__     ,atol=1e-5     ,rtol=1e-5							)
   @unittest.skip("Logits are not exactly the same, once we fix the instabalities somehow, will update!"							)
   @slow
   def     __UpperCamelCase						( self			:      str							):
     SCREAMING_SNAKE_CASE:Union[str, Any]						      =					[1, 306, 4_658, 278, 6_593, 310, 2_834, 338]
     SCREAMING_SNAKE_CASE:Union[str, Any]						      =					LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-13b-chat-hf"     ,device_map="auto"							)
     SCREAMING_SNAKE_CASE:Union[str, Any]						      =					model(torch.tensor(SCREAMING_SNAKE_CASE__							)							)
     # Expected mean on dim = -1
     SCREAMING_SNAKE_CASE:Optional[int]						      =					torch.tensor([[-0.8_562, -1.8_520, -0.7_551, -0.4_162, -1.5_161, -1.2_038, -2.4_823, -2.3_254]]							)
     torch.testing.assert_close(out.mean(-1							)     ,SCREAMING_SNAKE_CASE__     ,atol=1e-2     ,rtol=1e-2							)
     # slicing logits[0, 0, 0:30]
     # fmt: off
     SCREAMING_SNAKE_CASE:Optional[Any]						      =					torch.tensor([-2.2_227, 4.8_828, 0.9_023, -0.4_578, -0.7_871, -0.1_033, -0.6_221, -0.5_786, -0.7_803, -1.0_674, -1.2_920, -0.1_570, 0.8_008, 2.0_723, -0.9_497, 0.2_771, -2.2_227, -0.7_612, -1.4_346, -1.2_061, -1.6_426, -0.3_000, -0.7_139, -1.1_934, -1.8_691, -1.6_973, -1.5_947, -1.2_705, -0.3_523, -0.5_513]							)
     # fmt: on
     torch.testing.assert_close(out.mean(-1							)     ,SCREAMING_SNAKE_CASE__     ,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     __UpperCamelCase						( self			:      List[str]							):
     SCREAMING_SNAKE_CASE:int						      =					[1, 306, 4_658, 278, 6_593, 310, 2_834, 338]
     SCREAMING_SNAKE_CASE:Union[str, Any]						      =					LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-70b-hf"     ,device_map="auto"							)
     SCREAMING_SNAKE_CASE:Optional[int]						      =					model(torch.tensor(SCREAMING_SNAKE_CASE__							)							)

     SCREAMING_SNAKE_CASE:int						      =					torch.tensor(
         [[-4.2_327, -3.3_360, -4.6_665, -4.7_631, -1.8_180, -3.4_170, -1.4_211, -3.1_810]]     ,dtype=torch.floataa							)
     torch.testing.assert_close(out.mean(-1							)     ,SCREAMING_SNAKE_CASE__     ,atol=1e-2     ,rtol=1e-2							)
     # fmt: off
     SCREAMING_SNAKE_CASE:int						      =					torch.tensor([-9.4_922, -3.9_551, 1.7_998, -5.6_758, -5.1_055, -5.8_984, -4.8_320, -6.8_086, -6.5_391, -5.6_172, -5.5_820, -5.5_352, 1.7_881, 3.6_289, -6.5_117, -3.4_785, -9.5_000, -6.0_352, -6.8_125, -6.0_195, -6.6_836, -5.4_727, -6.2_812, -6.0_391, -7.3_398, -7.4_297, -7.4_844, -6.5_820, -5.8_789, -5.5_312]							)
     # fmt: on
     torch.testing.assert_close(out[0, 0, :30]     ,SCREAMING_SNAKE_CASE__     ,atol=1e-5     ,rtol=1e-5							)







   @unittest.skip("Model is curently gated"							)
   @slow
   def     __UpperCamelCase						( self			:      Optional[Any]							):
     SCREAMING_SNAKE_CASE:Union[str, 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"
     SCREAMING_SNAKE_CASE:Union[str, Any]						      =					"Simply put, the theory of relativity states that "
     SCREAMING_SNAKE_CASE:Union[str, Any]						      =					LlamaTokenizer.from_pretrained("meta-llama/Llama-2-13b-chat-hf"							)
     SCREAMING_SNAKE_CASE:List[str]						      =					tokenizer.encode(SCREAMING_SNAKE_CASE__     ,return_tensors="pt"							)
     SCREAMING_SNAKE_CASE:Optional[Any]						      =					LlamaForCausalLM.from_pretrained(
         "meta-llama/Llama-2-13b-chat-hf"     ,device_map="sequential"     ,use_safetensors=SCREAMING_SNAKE_CASE__							)

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









'''simple docstring'''

import argparse
import random

import joblib
import numpy as np
import torch
from igf.igf import (
    SecondaryLearner,
    collect_objective_set,
    compute_perplexity,
    generate_datasets,
    load_gpta,
    recopy_gpta,
    set_seed,
    train_secondary_learner,
)
from torch.utils.data import DataLoader, RandomSampler

from transformers import GPTaLMHeadModel





def    A_  ( snake_case=32      ,       snake_case=10      ,       snake_case=100      ,       snake_case=1026      ,       snake_case=True      ,       snake_case="data/tokenized_stories_train_wikitext103.jbl"      ,       snake_case="igf_context_pairs.jbl"      ,       ):
  set_seed(3 )
  # generate train_data and objective_set
  SCREAMING_SNAKE_CASE      ,    SCREAMING_SNAKE_CASE:Optional[int]						      =					generate_datasets(
      snake_case      ,       snake_case      ,       number=snake_case      ,       min_len=1026      ,       trim=snake_case )
  # keeps model same across runs
  set_seed(4 )
  # model, lm_optimizer, lm_scheduler = recopy_gpt2(model, device, max_steps) # store original model weights
  # can we train on GPU?
  SCREAMING_SNAKE_CASE:Tuple						      =					torch.device("cuda:0" if torch.cuda.is_available() else "cpu" )

  # load pretrained model
  SCREAMING_SNAKE_CASE:Optional[Any]						      =					load_gpta("gpt2" ).to(snake_case )
  print("computing perplexity on objective set" )
  SCREAMING_SNAKE_CASE:str						      =					compute_perplexity(snake_case      ,       snake_case      ,       snake_case ).item()
  print("perplexity on objective set:"      ,       snake_case )

  # collect igf pairs and save to file demo.jbl
  collect_objective_set(snake_case      ,       snake_case      ,       snake_case      ,       snake_case      ,       snake_case      ,       snake_case      ,       snake_case      ,       snake_case )

  # clean up, delete model and data we don't need anymore
  del model, train_data, objective_set
  torch.cuda.empty_cache()





def    A_  ( snake_case      ,       snake_case=15      ,       snake_case=128      ,       snake_case=100      ,       snake_case="igf_model.pt"      ,       ):
  set_seed(42 )

  # Load pre-trained model
  SCREAMING_SNAKE_CASE:int						      =					GPTaLMHeadModel.from_pretrained("gpt2" )

  # Initialize secondary learner to use embedding weights of model
  SCREAMING_SNAKE_CASE:Optional[Any]						      =					SecondaryLearner(snake_case )

  # Train secondary learner
  SCREAMING_SNAKE_CASE:int						      =					train_secondary_learner(
      snake_case      ,       snake_case      ,       max_epochs=snake_case      ,       batch_size=snake_case      ,       eval_freq=100      ,       igf_model_path=snake_case      ,       )

  del model, secondary_learner_train_data
  torch.cuda.empty_cache()

  return secondary_learner





def    A_  ( snake_case      ,       snake_case      ,       snake_case      ,       snake_case=32      ,       snake_case=1000      ,       snake_case=16      ,       snake_case=1.0      ,       snake_case=recopy_gpta      ,       snake_case=None      ,       snake_case=10      ,       snake_case="gpt2_finetuned.pt"      ,       ):
  SCREAMING_SNAKE_CASE:str						      =					torch.device("cuda:0" if torch.cuda.is_available() else "cpu" )
  SCREAMING_SNAKE_CASE:Dict						      =					RandomSampler(snake_case )
  SCREAMING_SNAKE_CASE:str						      =					DataLoader(snake_case      ,       sampler=snake_case )

  SCREAMING_SNAKE_CASE:str						      =					max_steps // (len(snake_case )) + 1
  SCREAMING_SNAKE_CASE:Tuple						      =					0
  SCREAMING_SNAKE_CASE:List[Any]						      =					torch.zeros((1, context_len)      ,       dtype=torch.long      ,       device=snake_case )
  SCREAMING_SNAKE_CASE      ,    SCREAMING_SNAKE_CASE      ,    SCREAMING_SNAKE_CASE:int						      =					recopy_model(snake_case      ,       snake_case      ,       snake_case )

  model.train()
  if secondary_learner is not None:
    secondary_learner.to(snake_case )
    secondary_learner.eval()
  SCREAMING_SNAKE_CASE:List[Any]						      =					[]
  SCREAMING_SNAKE_CASE:Tuple						      =					0

  SCREAMING_SNAKE_CASE:Tuple						      =					[]
  SCREAMING_SNAKE_CASE:Tuple						      =					[]

  # Compute the performance of the transformer model at the beginning
  SCREAMING_SNAKE_CASE:int						      =					compute_perplexity(snake_case      ,       snake_case      ,       snake_case )
  test_perps.append(snake_case )
  print("Test perplexity, step"      ,       snake_case      ,       ":"      ,       snake_case )
  for epoch in range(int(snake_case ) ):
    for step, example in enumerate(snake_case ):
      torch.cuda.empty_cache()
      SCREAMING_SNAKE_CASE:str						      =					random.randint(0      ,       example.size(2 ) - context_len - 1 )
      SCREAMING_SNAKE_CASE:List[str]						      =					example[0, 0, start : start + context_len]
      lm_optimizer.zero_grad()
      SCREAMING_SNAKE_CASE:Any						      =					model(snake_case      ,       labels=snake_case )
      SCREAMING_SNAKE_CASE:List[Any]						      =					True

      if secondary_learner is not None:
        SCREAMING_SNAKE_CASE:List[str]						      =					secondary_learner.forward(
            torch.tensor(snake_case      ,       dtype=torch.long      ,       device=snake_case ).unsqueeze(0 ) )[0].item()
        observed_qs.append(float(snake_case ) )

        # Here we implement the simple non-constant threshold for the predicted IG(X) value
        # We will decay the selectivity of our secondary learner filter from
        # 1 standard deviation above average to 1 below average after 10 batches.

        if global_step == 10:
          SCREAMING_SNAKE_CASE:int						      =					-1
        if predicted_q < threshold:
          SCREAMING_SNAKE_CASE:Optional[int]						      =					False

            # If we passed the filter, add the context to the batch!
      if do_backprop:
        contexts.append(np.array(context.cpu() ) )
        SCREAMING_SNAKE_CASE:Union[str, Any]						      =					outputs[0]
        lm_loss.backward()
        examples += 1

      del outputs

      # Once the batch is filled with enough contexts, backprop on the batch.
      if examples == batch_size:
        torch.cuda.empty_cache()
        SCREAMING_SNAKE_CASE:int						      =					0
        # Do LM backprop
        torch.nn.utils.clip_grad_norm_(model.parameters()      ,       3.0 )
        lm_optimizer.step()
        lm_scheduler.step()  # Update learning rate schedule
        global_step += 1
        # Compute the performance of the transformer model at this batch
        if global_step % eval_interval == 0:
          SCREAMING_SNAKE_CASE:Optional[Any]						      =					compute_perplexity(snake_case      ,       snake_case      ,       snake_case )
          test_perps.append(snake_case )

          print("Test perplexity, step"      ,       snake_case      ,       ":"      ,       snake_case )
            # Break out of the loop after 60 batches
      if max_steps > 0 and global_step > 60:
        break
    if max_steps > 0 and global_step > 60:
      break

    # save finetuned transformer model
  torch.save(model.state_dict()      ,       snake_case )
  torch.cuda.empty_cache()
  # Do some cleaning up so we can reinitialize for the next run of this function
  del lm_optimizer
  del lm_scheduler
  return model





def    A_  ( ):
  SCREAMING_SNAKE_CASE:Optional[int]						      =					argparse.ArgumentParser(description="Fine-tune a transformer model with IGF on a language modeling task" )

  # Required parameters
  parser.add_argument(
      "--data_dir"      ,       default=snake_case      ,       type=snake_case      ,       required=snake_case      ,       help="The input data dir. Should contain data files for WikiText."      ,       )
  parser.add_argument(
      "--model_name_or_path"      ,       default=snake_case      ,       type=snake_case      ,       required=snake_case      ,       help="Path to pretrained model or model identifier from huggingface.co/models"      ,       )
  parser.add_argument(
      "--data_file"      ,       type=snake_case      ,       default=snake_case      ,       help=(
          "A jbl file containing tokenized data which can be split as objective dataset, "
          "train_dataset and test_dataset."
      )      ,       )

  parser.add_argument(
      "--igf_data_file"      ,       type=snake_case      ,       default=snake_case      ,       help="A jbl file containing the context and information gain pairs to train secondary learner."      ,       )

  parser.add_argument(
      "--output_dir"      ,       default=snake_case      ,       type=snake_case      ,       required=snake_case      ,       help="The output directory where the final fine-tuned model is stored."      ,       )

  parser.add_argument(
      "--tokenizer_name"      ,       default=snake_case      ,       type=snake_case      ,       help="Pretrained tokenizer name or path if not the same as model_name"      ,       )
  parser.add_argument("--seed"      ,       type=snake_case      ,       default=snake_case      ,       help="A seed for reproducible training." )

  parser.add_argument(
      "--context_len"      ,       default=32      ,       type=snake_case      ,       help=(
          "The maximum total input sequence length after tokenization. Sequences longer "
          "than this will be truncated, sequences shorter will be padded."
      )      ,       )

  parser.add_argument(
      "--size_objective_set"      ,       default=100      ,       type=snake_case      ,       help="number of articles that are long enough to be used as our objective set"      ,       )
  parser.add_argument(
      "--eval_freq"      ,       default=100      ,       type=snake_case      ,       help="secondary model evaluation is triggered at eval_freq" )

  parser.add_argument("--max_steps"      ,       default=1000      ,       type=snake_case      ,       help="To calculate training epochs" )

  parser.add_argument(
      "--secondary_learner_batch_size"      ,       default=128      ,       type=snake_case      ,       help="batch size of training data for secondary learner"      ,       )

  parser.add_argument(
      "--batch_size"      ,       default=16      ,       type=snake_case      ,       help="batch size of training data of language model(gpt2) " )

  parser.add_argument(
      "--eval_interval"      ,       default=10      ,       type=snake_case      ,       help=(
          "decay the selectivity of our secondary learner filter from"
          "1 standard deviation above average to 1 below average after 10 batches"
      )      ,       )

  parser.add_argument(
      "--number"      ,       default=100      ,       type=snake_case      ,       help="The number of examples split to be used as objective_set/test_data" )

  parser.add_argument(
      "--min_len"      ,       default=1026      ,       type=snake_case      ,       help="The minimum length of the article to be used as objective set" )

  parser.add_argument(
      "--secondary_learner_max_epochs"      ,       default=15      ,       type=snake_case      ,       help="number of epochs to train secondary learner" )

  parser.add_argument("--trim"      ,       default=snake_case      ,       type=snake_case      ,       help="truncate the example if it exceeds context length" )

  parser.add_argument(
      "--threshold"      ,       default=1.0      ,       type=snake_case      ,       help=(
          "The threshold value used by secondary learner to filter the train_data and allow only"
          " informative data as input to the model"
      )      ,       )

  parser.add_argument("--finetuned_model_name"      ,       default="gpt2_finetuned.pt"      ,       type=snake_case      ,       help="finetuned_model_name" )

  parser.add_argument(
      "--recopy_model"      ,       default=snake_case      ,       type=snake_case      ,       help="Reset the model to the original pretrained GPT-2 weights after each iteration"      ,       )

  # function calls
  # Collecting *n* pairs of context and information gain(X, IG(X)) for training the secondary learner
  generate_n_pairs(
      context_len=32      ,       max_steps=10      ,       size_objective_set=100      ,       min_len=1026      ,       trim=snake_case      ,       data_file="data/tokenized_stories_train_wikitext103.jbl"      ,       igf_data_file="igf_context_pairs.jbl"      ,       )

  # Load train data for secondary learner
  SCREAMING_SNAKE_CASE:Union[str, Any]						      =					joblib.load("data/IGF_values.jbl" )

  # Train secondary learner
  SCREAMING_SNAKE_CASE:Optional[int]						      =					training_secondary_learner(
      snake_case      ,       secondary_learner_max_epochs=15      ,       secondary_learner_batch_size=128      ,       eval_freq=100      ,       igf_model_path="igf_model.pt"      ,       )

  # load pretrained gpt2 model
  SCREAMING_SNAKE_CASE:Dict						      =					GPTaLMHeadModel.from_pretrained("gpt2" )
  set_seed(42 )

  # Generate train and test data to train and evaluate gpt2 model
  SCREAMING_SNAKE_CASE      ,    SCREAMING_SNAKE_CASE:Optional[Any]						      =					generate_datasets(
      context_len=32      ,       file="data/tokenized_stories_train_wikitext103.jbl"      ,       number=100      ,       min_len=1026      ,       trim=snake_case )

  # fine-tuning of the gpt2 model using igf (Information Gain Filtration)
  finetune(
      snake_case      ,       snake_case      ,       snake_case      ,       context_len=32      ,       max_steps=1000      ,       batch_size=16      ,       threshold=1.0      ,       recopy_model=snake_case      ,       secondary_learner=snake_case      ,       eval_interval=10      ,       finetuned_model_name="gpt2_finetuned.pt"      ,       )


if __name__ == "__main__":
    main()






