


"""simple docstring"""





import argparse
import glob
import logging
import os
import time
from argparse import Namespace

import numpy as np
import torch
from lightning_base import BaseTransformer, add_generic_args, generic_train
from torch.utils.data import DataLoader, TensorDataset

from transformers import glue_compute_metrics as compute_metrics
from transformers import glue_convert_examples_to_features as convert_examples_to_features
from transformers import glue_output_modes, glue_tasks_num_labels
from transformers import glue_processors as processors


snake_case__		:     Any    				=				logging.getLogger(__name__)



class 		snake_case_(       __A			):
					__UpperCamelCase  =					'''sequence-classification'''






					def __init__(	self  :					Any					,				UpperCamelCase_  :					Optional[int]				):

											if type(UpperCamelCase_				) == dict:
																	lowerCAmelCase					:	Any       =   Namespace(**UpperCamelCase_				)
											lowerCAmelCase					:	int       =   glue_output_modes[hparams.task]
											lowerCAmelCase					:	Any       =   glue_tasks_num_labels[hparams.task]

											super().__init__(UpperCamelCase_					,				UpperCamelCase_					,				self.mode				)






					def      lowerCamelCase__     (	self  :					Dict					,				**UpperCamelCase_  :					Any				):
											return self.model(**UpperCamelCase_				)






					def      lowerCamelCase__     (	self  :					Optional[int]					,				UpperCamelCase_  :					Union[str, Any]					,				UpperCamelCase_  :					Optional[Any]				):
											lowerCAmelCase					:	Optional[int]       =   {'''input_ids''': batch[0], '''attention_mask''': batch[1], '''labels''': batch[3]}

											if self.config.model_type not in ["distilbert", "bart"]:
																	lowerCAmelCase					:	str       =   batch[2] if self.config.model_type in ['''bert''', '''xlnet''', '''albert'''] else None

											lowerCAmelCase					:	Union[str, Any]       =   self(**UpperCamelCase_				)
											lowerCAmelCase					:	List[Any]       =   outputs[0]

											lowerCAmelCase					:	str       =   self.trainer.lr_schedulers[0]['''scheduler''']
											lowerCAmelCase					:	Union[str, Any]       =   {'''loss''': loss, '''rate''': lr_scheduler.get_last_lr()[-1]}
											return {"loss": loss, "log": tensorboard_logs}






					def      lowerCamelCase__     (	self  :					Union[str, Any]				):
											lowerCAmelCase					:	Union[str, Any]       =   self.hparams
											lowerCAmelCase					:	str       =   processors[args.task]()
											lowerCAmelCase					:	Optional[int]       =   processor.get_labels()

											for mode in ["train", "dev"]:
																	lowerCAmelCase					:	List[str]       =   self._feature_file(UpperCamelCase_				)
																	if os.path.exists(UpperCamelCase_				) and not args.overwrite_cache:
																							logger.info('''Loading features from cached file %s'''					,				UpperCamelCase_				)
																	else:
																							logger.info('''Creating features from dataset file at %s'''					,				args.data_dir				)
																							lowerCAmelCase					:	Optional[int]       =   (
																							    processor.get_dev_examples(args.data_dir				)
																							    if mode == '''dev'''
																							    else processor.get_train_examples(args.data_dir				)
																							)
																							lowerCAmelCase					:	Optional[int]       =   convert_examples_to_features(
																							    UpperCamelCase_					,				self.tokenizer					,				max_length=args.max_seq_length					,				label_list=self.labels					,				output_mode=args.glue_output_mode					,				)
																							logger.info('''Saving features into cached file %s'''					,				UpperCamelCase_				)
																							torch.save(UpperCamelCase_					,				UpperCamelCase_				)






					def      lowerCamelCase__     (	self  :					List[str]					,				UpperCamelCase_  :					str					,				UpperCamelCase_  :					int					,				UpperCamelCase_  :					bool = False				):
											lowerCAmelCase					:	Tuple       =   '''dev''' if mode == '''test''' else mode

											lowerCAmelCase					:	Any       =   self._feature_file(UpperCamelCase_				)
											logger.info('''Loading features from cached file %s'''					,				UpperCamelCase_				)
											lowerCAmelCase					:	Optional[Any]       =   torch.load(UpperCamelCase_				)
											lowerCAmelCase					:	int       =   torch.tensor([f.input_ids for f in features]					,				dtype=torch.long				)
											lowerCAmelCase					:	str       =   torch.tensor([f.attention_mask for f in features]					,				dtype=torch.long				)
											lowerCAmelCase					:	Tuple       =   torch.tensor([f.token_type_ids for f in features]					,				dtype=torch.long				)
											if self.hparams.glue_output_mode == "classification":
																	lowerCAmelCase					:	int       =   torch.tensor([f.label for f in features]					,				dtype=torch.long				)
											elif self.hparams.glue_output_mode == "regression":
																	lowerCAmelCase					:	int       =   torch.tensor([f.label for f in features]					,				dtype=torch.float				)

											return DataLoader(
											    TensorDataset(UpperCamelCase_					,				UpperCamelCase_					,				UpperCamelCase_					,				UpperCamelCase_				)					,				batch_size=UpperCamelCase_					,				shuffle=UpperCamelCase_					,				)






					def      lowerCamelCase__     (	self  :					int					,				UpperCamelCase_  :					List[str]					,				UpperCamelCase_  :					Dict				):
											lowerCAmelCase					:	Any       =   {'''input_ids''': batch[0], '''attention_mask''': batch[1], '''labels''': batch[3]}

											if self.config.model_type not in ["distilbert", "bart"]:
																	lowerCAmelCase					:	Any       =   batch[2] if self.config.model_type in ['''bert''', '''xlnet''', '''albert'''] else None

											lowerCAmelCase					:	Union[str, Any]       =   self(**UpperCamelCase_				)
											lowerCAmelCase					:	List[Any]       =   outputs[:2]
											lowerCAmelCase					:	Union[str, Any]       =   logits.detach().cpu().numpy()
											lowerCAmelCase					:	int       =   inputs['''labels'''].detach().cpu().numpy()

											return {"val_loss": tmp_eval_loss.detach().cpu(), "pred": preds, "target": out_label_ids}






					def      lowerCamelCase__     (	self  :					List[str]					,				UpperCamelCase_  :					Dict				):
											lowerCAmelCase					:	Optional[Any]       =   torch.stack([x['''val_loss'''] for x in outputs]				).mean().detach().cpu().item()
											lowerCAmelCase					:	List[str]       =   np.concatenate([x['''pred'''] for x in outputs]					,				axis=0				)

											if self.hparams.glue_output_mode == "classification":
																	lowerCAmelCase					:	Optional[Any]       =   np.argmax(UpperCamelCase_					,				axis=1				)
											elif self.hparams.glue_output_mode == "regression":
																	lowerCAmelCase					:	int       =   np.squeeze(UpperCamelCase_				)

											lowerCAmelCase					:	Any       =   np.concatenate([x['''target'''] for x in outputs]					,				axis=0				)
											lowerCAmelCase					:	int       =   [[] for _ in range(out_label_ids.shape[0]				)]
											lowerCAmelCase					:	Tuple       =   [[] for _ in range(out_label_ids.shape[0]				)]

											lowerCAmelCase					:	Any       =   {**{'''val_loss''': val_loss_mean}, **compute_metrics(self.hparams.task					,				UpperCamelCase_					,				UpperCamelCase_				)}

											lowerCAmelCase					:	Optional[Any]       =   dict(results.items()				)
											lowerCAmelCase					:	str       =   results
											return ret, preds_list, out_label_list






					def      lowerCamelCase__     (	self  :					Optional[Any]					,				UpperCamelCase_  :					list				):
											lowerCAmelCase					:	Dict       =   self._eval_end(UpperCamelCase_				)
											lowerCAmelCase					:	Any       =   ret['''log''']
											return {"val_loss": logs["val_loss"], "log": logs, "progress_bar": logs}






					def      lowerCamelCase__     (	self  :					int					,				UpperCamelCase_  :					Tuple				):
											lowerCAmelCase					:	Union[str, Any]       =   self._eval_end(UpperCamelCase_				)
											lowerCAmelCase					:	Optional[Any]       =   ret['''log''']
											# `val_loss` is the key returned by `self._eval_end()` but actually refers to `test_loss`
											return {"avg_test_loss": logs["val_loss"], "log": logs, "progress_bar": logs}






					@staticmethod
					def      lowerCamelCase__     (	UpperCamelCase_  :					str					,				UpperCamelCase_  :					Any				):
											BaseTransformer.add_model_specific_args(UpperCamelCase_					,				UpperCamelCase_				)
											parser.add_argument(
											    '''--max_seq_length'''					,				default=1_2_8					,				type=UpperCamelCase_					,				help=(
											        '''The maximum total input sequence length after tokenization. Sequences longer '''
											        '''than this will be truncated, sequences shorter will be padded.'''
											    )					,				)

											parser.add_argument(
											    '''--task'''					,				default=''''''					,				type=UpperCamelCase_					,				required=UpperCamelCase_					,				help='''The GLUE task to run'''					,				)
											parser.add_argument(
											    '''--gpus'''					,				default=0					,				type=UpperCamelCase_					,				help='''The number of GPUs allocated for this, it is by default 0 meaning none'''					,				)

											parser.add_argument(
											    '''--overwrite_cache'''					,				action='''store_true'''					,				help='''Overwrite the cached training and evaluation sets'''				)

											return parser



def      _snake_case						(						):
						lowerCAmelCase					:	Tuple       =   argparse.ArgumentParser()
						add_generic_args(__A ,   os.getcwd()			)
						lowerCAmelCase					:	Any       =   GLUETransformer.add_model_specific_args(__A ,   os.getcwd()			)
						lowerCAmelCase					:	Optional[Any]       =   parser.parse_args()

						# If output_dir not provided, a folder will be generated in pwd
						if args.output_dir is None:
												lowerCAmelCase					:	Tuple       =   os.path.join(
												    '''./results''' ,   f'''{args.task}_{time.strftime("%Y%m%d_%H%M%S"			)}''' ,   )
												os.makedirs(args.output_dir			)

						lowerCAmelCase					:	str       =   GLUETransformer(__A			)
						lowerCAmelCase					:	Any       =   generic_train(__A ,   __A			)

						# Optionally, predict on dev set and write to output_dir
						if args.do_predict:
												lowerCAmelCase					:	Tuple       =   sorted(glob.glob(os.path.join(args.output_dir ,   '''checkpoint-epoch=*.ckpt'''			) ,   recursive=__A			)			)
												lowerCAmelCase					:	int       =   model.load_from_checkpoint(checkpoints[-1]			)
												return trainer.test(__A			)


if __name__ == "__main__":
	main()






"""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  :					int					,				UpperCamelCase_  :					List[Any]					,				UpperCamelCase_  :					Dict=1_3					,				UpperCamelCase_  :					Optional[Any]=7					,				UpperCamelCase_  :					Optional[int]=True					,				UpperCamelCase_  :					Dict=True					,				UpperCamelCase_  :					Optional[int]=False					,				UpperCamelCase_  :					Any=True					,				UpperCamelCase_  :					List[str]=9_9					,				UpperCamelCase_  :					Tuple=3_2					,				UpperCamelCase_  :					Optional[Any]=5					,				UpperCamelCase_  :					str=4					,				UpperCamelCase_  :					Any=3_7					,				UpperCamelCase_  :					Optional[Any]="gelu"					,				UpperCamelCase_  :					Tuple=0.1					,				UpperCamelCase_  :					Union[str, Any]=0.1					,				UpperCamelCase_  :					Union[str, Any]=5_1_2					,				UpperCamelCase_  :					Union[str, Any]=1_6					,				UpperCamelCase_  :					Any=2					,				UpperCamelCase_  :					Optional[Any]=0.02					,				UpperCamelCase_  :					List[Any]=3					,				UpperCamelCase_  :					Any=4					,				UpperCamelCase_  :					int=None					,				):
           lowerCAmelCase					:	Any       =   parent
           lowerCAmelCase					:	Any       =   batch_size
           lowerCAmelCase					:	List[Any]       =   seq_length
           lowerCAmelCase					:	str       =   is_training
           lowerCAmelCase					:	List[Any]       =   use_input_mask
           lowerCAmelCase					:	Optional[int]       =   use_token_type_ids
           lowerCAmelCase					:	Union[str, Any]       =   use_labels
           lowerCAmelCase					:	List[str]       =   vocab_size
           lowerCAmelCase					:	Tuple       =   hidden_size
           lowerCAmelCase					:	int       =   num_hidden_layers
           lowerCAmelCase					:	Union[str, Any]       =   num_attention_heads
           lowerCAmelCase					:	Optional[int]       =   intermediate_size
           lowerCAmelCase					:	List[Any]       =   hidden_act
           lowerCAmelCase					:	int       =   hidden_dropout_prob
           lowerCAmelCase					:	Tuple       =   attention_probs_dropout_prob
           lowerCAmelCase					:	Optional[Any]       =   max_position_embeddings
           lowerCAmelCase					:	Optional[int]       =   type_vocab_size
           lowerCAmelCase					:	Tuple       =   type_sequence_label_size
           lowerCAmelCase					:	List[str]       =   initializer_range
           lowerCAmelCase					:	str       =   num_labels
           lowerCAmelCase					:	Optional[int]       =   num_choices
           lowerCAmelCase					:	Tuple       =   scope






     def      lowerCamelCase__     (	self  :					Optional[int]				):
           lowerCAmelCase					:	List[str]       =   ids_tensor([self.batch_size, self.seq_length]					,				self.vocab_size				)

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

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

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

           lowerCAmelCase					:	Optional[Any]       =   self.get_config()

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






     def      lowerCamelCase__     (	self  :					Tuple				):
           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      lowerCamelCase__     (	self  :					int					,				UpperCamelCase_  :					Any					,				UpperCamelCase_  :					Dict					,				UpperCamelCase_  :					Any					,				UpperCamelCase_  :					Dict					,				UpperCamelCase_  :					Union[str, Any]					,				UpperCamelCase_  :					List[str]					,				UpperCamelCase_  :					Tuple				):
           lowerCAmelCase					:	List[Any]       =   LlamaModel(config=UpperCamelCase_				)
           model.to(UpperCamelCase_				)
           model.eval()
           lowerCAmelCase					:	Dict       =   model(UpperCamelCase_					,				attention_mask=UpperCamelCase_				)
           lowerCAmelCase					:	Optional[int]       =   model(UpperCamelCase_				)
           self.parent.assertEqual(result.last_hidden_state.shape					,				(self.batch_size, self.seq_length, self.hidden_size)				)






     def      lowerCamelCase__     (	self  :					Any					,				UpperCamelCase_  :					Optional[Any]					,				UpperCamelCase_  :					Tuple					,				UpperCamelCase_  :					Tuple					,				UpperCamelCase_  :					Dict					,				UpperCamelCase_  :					int					,				UpperCamelCase_  :					Dict					,				UpperCamelCase_  :					Tuple					,				UpperCamelCase_  :					int					,				UpperCamelCase_  :					Any					,				):
           lowerCAmelCase					:	Tuple       =   True
           lowerCAmelCase					:	Optional[int]       =   LlamaModel(UpperCamelCase_				)
           model.to(UpperCamelCase_				)
           model.eval()
           lowerCAmelCase					:	List[Any]       =   model(
               UpperCamelCase_					,				attention_mask=UpperCamelCase_					,				encoder_hidden_states=UpperCamelCase_					,				encoder_attention_mask=UpperCamelCase_					,				)
           lowerCAmelCase					:	Dict       =   model(
               UpperCamelCase_					,				attention_mask=UpperCamelCase_					,				encoder_hidden_states=UpperCamelCase_					,				)
           lowerCAmelCase					:	Tuple       =   model(UpperCamelCase_					,				attention_mask=UpperCamelCase_				)
           self.parent.assertEqual(result.last_hidden_state.shape					,				(self.batch_size, self.seq_length, self.hidden_size)				)






     def      lowerCamelCase__     (	self  :					Any					,				UpperCamelCase_  :					int					,				UpperCamelCase_  :					Tuple					,				UpperCamelCase_  :					Any					,				UpperCamelCase_  :					List[Any]					,				UpperCamelCase_  :					str					,				UpperCamelCase_  :					int					,				UpperCamelCase_  :					int					,				UpperCamelCase_  :					int					,				UpperCamelCase_  :					str					,				):
           lowerCAmelCase					:	Optional[Any]       =   LlamaForCausalLM(config=UpperCamelCase_				)
           model.to(UpperCamelCase_				)
           model.eval()
           lowerCAmelCase					:	List[str]       =   model(UpperCamelCase_					,				attention_mask=UpperCamelCase_					,				labels=UpperCamelCase_				)
           self.parent.assertEqual(result.logits.shape					,				(self.batch_size, self.seq_length, self.vocab_size)				)






     def      lowerCamelCase__     (	self  :					str					,				UpperCamelCase_  :					List[str]					,				UpperCamelCase_  :					int					,				UpperCamelCase_  :					Tuple					,				UpperCamelCase_  :					Optional[Any]					,				UpperCamelCase_  :					int					,				UpperCamelCase_  :					str					,				UpperCamelCase_  :					Dict					,				UpperCamelCase_  :					Optional[Any]					,				UpperCamelCase_  :					List[Any]					,				):
           lowerCAmelCase					:	Union[str, Any]       =   True
           lowerCAmelCase					:	str       =   True
           lowerCAmelCase					:	Tuple       =   LlamaForCausalLM(config=UpperCamelCase_				)
           model.to(UpperCamelCase_				)
           model.eval()

           # first forward pass
           lowerCAmelCase					:	Optional[Any]       =   model(
               UpperCamelCase_					,				attention_mask=UpperCamelCase_					,				encoder_hidden_states=UpperCamelCase_					,				encoder_attention_mask=UpperCamelCase_					,				use_cache=UpperCamelCase_					,				)
           lowerCAmelCase					:	Dict       =   outputs.past_key_values

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

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

           lowerCAmelCase					:	Dict       =   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					:	Any       =   output_from_no_past[:, -3:, random_slice_idx].detach()
           lowerCAmelCase					:	Optional[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(UpperCamelCase_					,				UpperCamelCase_					,				atol=1E-3				)				)






     def      lowerCamelCase__     (	self  :					Union[str, Any]				):
           lowerCAmelCase					:	Dict       =   self.prepare_config_and_inputs()
           (
               (
               lowerCAmelCase
           ),      (
               lowerCAmelCase
           ),      (
               lowerCAmelCase
           ),      (
               lowerCAmelCase
           ),      (
               lowerCAmelCase
           ),      (
               lowerCAmelCase
           ),      (
               lowerCAmelCase
           ),      
           )					:	Tuple       =   config_and_inputs
           lowerCAmelCase					:	Optional[int]       =   {'''input_ids''': input_ids, '''attention_mask''': input_mask}
           return config, inputs_dict



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






     def      lowerCamelCase__     (	self  :					Optional[Any]				):
           lowerCAmelCase					:	Any       =   LlamaModelTester(self				)
           lowerCAmelCase					:	Dict       =   ConfigTester(self					,				config_class=UpperCamelCase_					,				hidden_size=3_7				)






     def      lowerCamelCase__     (	self  :					str				):
           self.config_tester.run_common_tests()






     def      lowerCamelCase__     (	self  :					Tuple				):
           lowerCAmelCase					:	Tuple       =   self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*UpperCamelCase_				)






     def      lowerCamelCase__     (	self  :					List[Any]				):
           lowerCAmelCase					:	str       =   self.model_tester.prepare_config_and_inputs()
           for type in ["absolute", "relative_key", "relative_key_query"]:
                 lowerCAmelCase					:	str       =   type
                 self.model_tester.create_and_check_model(*UpperCamelCase_				)






     def      lowerCamelCase__     (	self  :					List[Any]				):
           lowerCAmelCase,      lowerCAmelCase					:	Tuple       =   self.model_tester.prepare_config_and_inputs_for_common()
           lowerCAmelCase					:	List[str]       =   3
           lowerCAmelCase					:	List[str]       =   input_dict['''input_ids''']
           lowerCAmelCase					:	List[str]       =   input_ids.ne(1				).to(UpperCamelCase_				)
           lowerCAmelCase					:	Tuple       =   ids_tensor([self.model_tester.batch_size]					,				self.model_tester.type_sequence_label_size				)
           lowerCAmelCase					:	Union[str, Any]       =   LlamaForSequenceClassification(UpperCamelCase_				)
           model.to(UpperCamelCase_				)
           model.eval()
           lowerCAmelCase					:	List[Any]       =   model(UpperCamelCase_					,				attention_mask=UpperCamelCase_					,				labels=UpperCamelCase_				)
           self.assertEqual(result.logits.shape					,				(self.model_tester.batch_size, self.model_tester.num_labels)				)






     def      lowerCamelCase__     (	self  :					Optional[int]				):
           lowerCAmelCase,      lowerCAmelCase					:	Optional[int]       =   self.model_tester.prepare_config_and_inputs_for_common()
           lowerCAmelCase					:	Any       =   3
           lowerCAmelCase					:	int       =   '''single_label_classification'''
           lowerCAmelCase					:	Tuple       =   input_dict['''input_ids''']
           lowerCAmelCase					:	Tuple       =   input_ids.ne(1				).to(UpperCamelCase_				)
           lowerCAmelCase					:	str       =   ids_tensor([self.model_tester.batch_size]					,				self.model_tester.type_sequence_label_size				)
           lowerCAmelCase					:	Tuple       =   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)				)






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






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

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






     @unittest.skip('''Logits are not exactly the same, once we fix the instabalities somehow, will update!'''				)
     @slow
     def      lowerCamelCase__     (	self  :					List[Any]				):
           lowerCAmelCase					:	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]
           lowerCAmelCase					:	Optional[Any]       =   LlamaForCausalLM.from_pretrained('''meta-llama/Llama-2-7b-hf'''					,				device_map='''auto'''				)
           lowerCAmelCase					:	str       =   model(torch.tensor([input_ids]				)				)
           # Expected mean on dim = -1
           lowerCAmelCase					:	int       =   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				)					,				UpperCamelCase_					,				atol=1E-2					,				rtol=1E-2				)
           # slicing logits[0, 0, 0:30]
           # fmt: off
           lowerCAmelCase					:	Tuple       =   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, :3_0]					,				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      lowerCamelCase__     (	self  :					Dict				):
           lowerCAmelCase					:	str       =   [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]
           lowerCAmelCase					:	Dict       =   LlamaForCausalLM.from_pretrained('''meta-llama/Llama-2-13b-hf'''					,				device_map='''auto'''				)
           lowerCAmelCase					:	str       =   model(torch.tensor(UpperCamelCase_				)				)
           # Expected mean on dim = -1
           lowerCAmelCase					:	Any       =   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				)					,				UpperCamelCase_					,				atol=1E-2					,				rtol=1E-2				)
           # slicing logits[0, 0, 0:30]
           # fmt: off
           lowerCAmelCase					:	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, :3_0]					,				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      lowerCamelCase__     (	self  :					Optional[int]				):
           lowerCAmelCase					:	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]
           lowerCAmelCase					:	List[str]       =   LlamaForCausalLM.from_pretrained('''meta-llama/Llama-2-13b-chat-hf'''					,				device_map='''auto'''				)
           lowerCAmelCase					:	List[Any]       =   model(torch.tensor(UpperCamelCase_				)				)
           # Expected mean on dim = -1
           lowerCAmelCase					:	List[str]       =   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				)					,				UpperCamelCase_					,				atol=1E-2					,				rtol=1E-2				)
           # slicing logits[0, 0, 0:30]
           # fmt: off
           lowerCAmelCase					:	Dict       =   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				)					,				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      lowerCamelCase__     (	self  :					List[str]				):
           lowerCAmelCase					:	Optional[Any]       =   [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]
           lowerCAmelCase					:	Optional[int]       =   LlamaForCausalLM.from_pretrained('''meta-llama/Llama-2-70b-hf'''					,				device_map='''auto'''				)
           lowerCAmelCase					:	Any       =   model(torch.tensor(UpperCamelCase_				)				)

           lowerCAmelCase					:	Optional[Any]       =   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				)					,				UpperCamelCase_					,				atol=1E-2					,				rtol=1E-2				)
           # fmt: off
           lowerCAmelCase					:	Any       =   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, :3_0]					,				UpperCamelCase_					,				atol=1E-5					,				rtol=1E-5				)






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

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


