


"""simple docstring"""





import argparse
import glob
import logging
import os
from argparse import Namespace
from importlib import import_module

import numpy as np
import torch
from lightning_base import BaseTransformer, add_generic_args, generic_train
from seqeval.metrics import accuracy_score, fa_score, precision_score, recall_score
from torch.nn import CrossEntropyLoss
from torch.utils.data import DataLoader, TensorDataset
from utils_ner import TokenClassificationTask


snake_case__		:     int    				=				logging.getLogger(__name__)



class 		snake_case_(       A_			):
					__UpperCamelCase  =					'''token-classification'''






					def __init__(	self  :					Tuple					,				UpperCamelCase_  :					Any				):
											if type(UpperCamelCase_				) == dict:
																	lowerCAmelCase					:	Any       =   Namespace(**UpperCamelCase_				)
											lowerCAmelCase					:	str       =   import_module('''tasks'''				)
											try:
																	lowerCAmelCase					:	Optional[Any]       =   getattr(UpperCamelCase_					,				hparams.task_type				)
																	lowerCAmelCase					:	TokenClassificationTask       =   token_classification_task_clazz()
											except AttributeError:
																	raise ValueError(
																	    F'''Task {hparams.task_type} needs to be defined as a TokenClassificationTask subclass in {module}. '''
																	    F'''Available tasks classes are: {TokenClassificationTask.__subclasses__()}'''				)
											lowerCAmelCase					:	Any       =   self.token_classification_task.get_labels(hparams.labels				)
											lowerCAmelCase					:	Tuple       =   CrossEntropyLoss().ignore_index
											super().__init__(UpperCamelCase_					,				len(self.labels				)					,				self.mode				)






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






					def      lowerCamelCase__     (	self  :					Union[str, Any]					,				UpperCamelCase_  :					Tuple					,				UpperCamelCase_  :					int				):
											lowerCAmelCase					:	Tuple       =   {"""input_ids""": batch[0], """attention_mask""": batch[1], """labels""": batch[3]}
											if self.config.model_type != "distilbert":
																	lowerCAmelCase					:	Union[str, Any]       =   (
																	    batch[2] if self.config.model_type in ["""bert""", """xlnet"""] else None
																	)  # XLM and RoBERTa don"t use token_type_ids

											lowerCAmelCase					:	List[str]       =   self(**UpperCamelCase_				)
											lowerCAmelCase					:	Dict       =   outputs[0]
											# tensorboard_logs = {"loss": loss, "rate": self.lr_scheduler.get_last_lr()[-1]}
											return {"loss": loss}






					def      lowerCamelCase__     (	self  :					List[str]				):
											lowerCAmelCase					:	Dict       =   self.hparams
											for mode in ["train", "dev", "test"]:
																	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_				)
																							lowerCAmelCase					:	Union[str, Any]       =   torch.load(UpperCamelCase_				)
																	else:
																							logger.info('''Creating features from dataset file at %s'''					,				args.data_dir				)
																							lowerCAmelCase					:	int       =   self.token_classification_task.read_examples_from_file(args.data_dir					,				UpperCamelCase_				)
																							lowerCAmelCase					:	Tuple       =   self.token_classification_task.convert_examples_to_features(
																							    UpperCamelCase_					,				self.labels					,				args.max_seq_length					,				self.tokenizer					,				cls_token_at_end=bool(self.config.model_type in ['''xlnet''']				)					,				cls_token=self.tokenizer.cls_token					,				cls_token_segment_id=2 if self.config.model_type in ['''xlnet'''] else 0					,				sep_token=self.tokenizer.sep_token					,				sep_token_extra=UpperCamelCase_					,				pad_on_left=bool(self.config.model_type in ['''xlnet''']				)					,				pad_token=self.tokenizer.pad_token_id					,				pad_token_segment_id=self.tokenizer.pad_token_type_id					,				pad_token_label_id=self.pad_token_label_id					,				)
																							logger.info('''Saving features into cached file %s'''					,				UpperCamelCase_				)
																							torch.save(UpperCamelCase_					,				UpperCamelCase_				)






					def      lowerCamelCase__     (	self  :					Union[str, Any]					,				UpperCamelCase_  :					Tuple					,				UpperCamelCase_  :					Tuple					,				UpperCamelCase_  :					Optional[Any] = False				):
											lowerCAmelCase					:	Any       =   self._feature_file(UpperCamelCase_				)
											logger.info('''Loading features from cached file %s'''					,				UpperCamelCase_				)
											lowerCAmelCase					:	Optional[Any]       =   torch.load(UpperCamelCase_				)
											lowerCAmelCase					:	Tuple       =   torch.tensor([f.input_ids for f in features]					,				dtype=torch.long				)
											lowerCAmelCase					:	Optional[Any]       =   torch.tensor([f.attention_mask for f in features]					,				dtype=torch.long				)
											if features[0].token_type_ids is not None:
																	lowerCAmelCase					:	str       =   torch.tensor([f.token_type_ids for f in features]					,				dtype=torch.long				)
											else:
																	lowerCAmelCase					:	int       =   torch.tensor([0 for f in features]					,				dtype=torch.long				)
																	# HACK(we will not use this anymore soon)
											lowerCAmelCase					:	Any       =   torch.tensor([f.label_ids for f in features]					,				dtype=torch.long				)
											return DataLoader(
											    TensorDataset(UpperCamelCase_					,				UpperCamelCase_					,				UpperCamelCase_					,				UpperCamelCase_				)					,				batch_size=UpperCamelCase_				)






					def      lowerCamelCase__     (	self  :					Union[str, Any]					,				UpperCamelCase_  :					Any					,				UpperCamelCase_  :					Union[str, Any]				):
											"""Compute validation""" ""
											lowerCAmelCase					:	Optional[int]       =   {"""input_ids""": batch[0], """attention_mask""": batch[1], """labels""": batch[3]}
											if self.config.model_type != "distilbert":
																	lowerCAmelCase					:	Tuple       =   (
																	    batch[2] if self.config.model_type in ["""bert""", """xlnet"""] else None
																	)  # XLM and RoBERTa don"t use token_type_ids
											lowerCAmelCase					:	str       =   self(**UpperCamelCase_				)
											lowerCAmelCase					:	List[Any]       =   outputs[:2]
											lowerCAmelCase					:	List[Any]       =   logits.detach().cpu().numpy()
											lowerCAmelCase					:	Dict       =   inputs["""labels"""].detach().cpu().numpy()
											return {"val_loss": tmp_eval_loss.detach().cpu(), "pred": preds, "target": out_label_ids}






					def      lowerCamelCase__     (	self  :					Tuple					,				UpperCamelCase_  :					int				):
											lowerCAmelCase					:	List[str]       =   torch.stack([x['''val_loss'''] for x in outputs]				).mean()
											lowerCAmelCase					:	Optional[int]       =   np.concatenate([x['''pred'''] for x in outputs]					,				axis=0				)
											lowerCAmelCase					:	List[str]       =   np.argmax(UpperCamelCase_					,				axis=2				)
											lowerCAmelCase					:	Optional[Any]       =   np.concatenate([x['''target'''] for x in outputs]					,				axis=0				)

											lowerCAmelCase					:	Optional[int]       =   dict(enumerate(self.labels				)				)
											lowerCAmelCase					:	List[str]       =   [[] for _ in range(out_label_ids.shape[0]				)]
											lowerCAmelCase					:	Any       =   [[] for _ in range(out_label_ids.shape[0]				)]

											for i in range(out_label_ids.shape[0]				):
																	for j in range(out_label_ids.shape[1]				):
																							if out_label_ids[i, j] != self.pad_token_label_id:
																													out_label_list[i].append(label_map[out_label_ids[i][j]]				)
																													preds_list[i].append(label_map[preds[i][j]]				)

											lowerCAmelCase					:	Tuple       =   {
											    """val_loss""": val_loss_mean,
											    """accuracy_score""": accuracy_score(UpperCamelCase_					,				UpperCamelCase_				),
											    """precision""": precision_score(UpperCamelCase_					,				UpperCamelCase_				),
											    """recall""": recall_score(UpperCamelCase_					,				UpperCamelCase_				),
											    """f1""": fa_score(UpperCamelCase_					,				UpperCamelCase_				),
											}

											lowerCAmelCase					:	Dict       =   dict(results.items()				)
											lowerCAmelCase					:	str       =   results
											return ret, preds_list, out_label_list






					def      lowerCamelCase__     (	self  :					int					,				UpperCamelCase_  :					Union[str, Any]				):
											# when stable
											lowerCAmelCase					:	str       =   self._eval_end(UpperCamelCase_				)
											lowerCAmelCase					:	Union[str, Any]       =   ret["""log"""]
											return {"val_loss": logs["val_loss"], "log": logs, "progress_bar": logs}






					def      lowerCamelCase__     (	self  :					int					,				UpperCamelCase_  :					Any				):
											# updating to test_epoch_end instead of deprecated test_end
											lowerCAmelCase					:	Union[str, Any]       =   self._eval_end(UpperCamelCase_				)

											# Converting to the dict required by pl
											# https://github.com/PyTorchLightning/pytorch-lightning/blob/master/\
											# pytorch_lightning/trainer/logging.py#L139
											lowerCAmelCase					:	List[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_  :					List[Any]					,				UpperCamelCase_  :					Optional[Any]				):
											# Add NER specific options
											BaseTransformer.add_model_specific_args(UpperCamelCase_					,				UpperCamelCase_				)
											parser.add_argument(
											    '''--task_type'''					,				default='''NER'''					,				type=UpperCamelCase_					,				help='''Task type to fine tune in training (e.g. NER, POS, etc)'''				)
											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(
											    '''--labels'''					,				default=''''''					,				type=UpperCamelCase_					,				help='''Path to a file containing all labels. If not specified, CoNLL-2003 labels are used.'''					,				)
											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


if __name__ == "__main__":
	snake_case__		:     Dict    				=				argparse.ArgumentParser()
	add_generic_args(parser, os.getcwd())
	snake_case__		:     Any    				=				NERTransformer.add_model_specific_args(parser, os.getcwd())
	snake_case__		:     List[Any]    				=				parser.parse_args()
	snake_case__		:     int    				=				NERTransformer(args)
	snake_case__		:     Any    				=				generic_train(model, args)

	if args.do_predict:
		# See https://github.com/huggingface/transformers/issues/3159
		# pl use this default format to create a checkpoint:
		# https://github.com/PyTorchLightning/pytorch-lightning/blob/master\
		# /pytorch_lightning/callbacks/model_checkpoint.py#L322
		snake_case__		:     int    				=				sorted(glob.glob(os.path.join(args.output_dir, '''checkpoint-epoch=*.ckpt'''), recursive=True))
		snake_case__		:     List[Any]    				=				model.load_from_checkpoint(checkpoints[-1])
		trainer.test(model)






"""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_				)


