


"""simple docstring"""





import sys
from pathlib import Path


snake_case__		:     str    				=				Path(__file__).resolve().parents[3] / '''src'''
sys.path.insert(1, str(git_repo_path))

import dataclasses  # noqa
import io  # noqa
import itertools  # noqa
import json  # noqa
import os  # noqa
import unittest  # noqa
from copy import deepcopy  # noqa

from parameterized import parameterized  # noqa
from transformers import TrainingArguments, is_torch_available  # noqa
from transformers.deepspeed import is_deepspeed_available  # noqa
from transformers.file_utils import WEIGHTS_NAME  # noqa
from transformers.testing_utils import (  # noqa
    CaptureLogger,
    ExtendSysPath,
    TestCasePlus,
    execute_subprocess_async,
    get_gpu_count,
    mockenv_context,
    require_deepspeed,
    require_torch_gpu,
    require_torch_multi_gpu,
    slow,
)
from transformers.trainer_utils import set_seed  # noqa


set_seed(42)

snake_case__		:     Union[str, Any]    				=				{'''base''': '''patrickvonplaten/wav2vec2_tiny_random''', '''robust''': '''patrickvonplaten/wav2vec2_tiny_random_robust'''}

snake_case__		:     str    				=				'''zero2'''
snake_case__		:     Optional[Any]    				=				'''zero3'''
snake_case__		:     Optional[int]    				=				[ZEROa, ZEROa]


def      _snake_case						(						_snake_case				:			Optional[Any] ,   _snake_case				:			int ,   _snake_case				:			str			):
						# customize the test name generator function as we want both params to appear in the sub-test
						# name, as by default it shows only the first param
						lowerCAmelCase					:	Union[str, Any]       =   parameterized.to_safe_name('''_'''.join(str(_snake_case			) for x in param.args			)			)
						return f'''{func.__name__}_{param_based_name}'''


# Cartesian-product of zero stages with models to test
snake_case__		:     str    				=				list(itertools.product(stages, models.keys()))






@slow
@require_deepspeed
@require_torch_gpu
class 		snake_case_(       a__			):






					@parameterized.expand(UpperCamelCase_					,				name_func=UpperCamelCase_				)
					def      lowerCamelCase__     (	self  :					List[str]					,				UpperCamelCase_  :					int					,				UpperCamelCase_  :					Dict				):
											self.run_and_check(
											    stage=UpperCamelCase_					,				model=UpperCamelCase_					,				distributed=UpperCamelCase_					,				fpaa=UpperCamelCase_					,				)






					@require_torch_multi_gpu
					@parameterized.expand(UpperCamelCase_					,				name_func=UpperCamelCase_				)
					def      lowerCamelCase__     (	self  :					Any					,				UpperCamelCase_  :					int					,				UpperCamelCase_  :					Optional[int]				):
											self.run_and_check(
											    stage=UpperCamelCase_					,				model=UpperCamelCase_					,				distributed=UpperCamelCase_					,				fpaa=UpperCamelCase_					,				)






					@parameterized.expand(UpperCamelCase_					,				name_func=UpperCamelCase_				)
					def      lowerCamelCase__     (	self  :					Union[str, Any]					,				UpperCamelCase_  :					Optional[int]					,				UpperCamelCase_  :					Optional[int]				):
											self.run_and_check(
											    stage=UpperCamelCase_					,				model=UpperCamelCase_					,				distributed=UpperCamelCase_					,				fpaa=UpperCamelCase_					,				)






					@require_torch_multi_gpu
					@parameterized.expand(UpperCamelCase_					,				name_func=UpperCamelCase_				)
					def      lowerCamelCase__     (	self  :					str					,				UpperCamelCase_  :					Tuple					,				UpperCamelCase_  :					List[str]				):
											self.run_and_check(
											    stage=UpperCamelCase_					,				model=UpperCamelCase_					,				distributed=UpperCamelCase_					,				fpaa=UpperCamelCase_					,				)






					def      lowerCamelCase__     (	self  :					int					,				UpperCamelCase_  :					int				):
											# XXX: run_asr is premature and doesn't save any results
											# so all we check for now is that the process didn't fail
											pass






					def      lowerCamelCase__     (	self  :					List[str]					,				UpperCamelCase_  :					str					,				UpperCamelCase_  :					str					,				UpperCamelCase_  :					int = 1_0					,				UpperCamelCase_  :					bool = True					,				UpperCamelCase_  :					bool = True					,				UpperCamelCase_  :					bool = True					,				):
											lowerCAmelCase					:	Tuple       =   models[model]

											lowerCAmelCase					:	Any       =   self.run_trainer(
											    stage=UpperCamelCase_					,				model_name=UpperCamelCase_					,				eval_steps=UpperCamelCase_					,				num_train_epochs=1					,				distributed=UpperCamelCase_					,				fpaa=UpperCamelCase_					,				)

											self.do_checks(UpperCamelCase_				)

											return output_dir






					def      lowerCamelCase__     (	self  :					Optional[Any]					,				UpperCamelCase_  :					str					,				UpperCamelCase_  :					str					,				UpperCamelCase_  :					int = 1_0					,				UpperCamelCase_  :					int = 1					,				UpperCamelCase_  :					bool = True					,				UpperCamelCase_  :					bool = True					,				):
											lowerCAmelCase					:	Any       =   self.get_auto_remove_tmp_dir('''./xxx'''					,				after=UpperCamelCase_				)
											lowerCAmelCase					:	Tuple       =   F'''
            --model_name_or_path {model_name}
            --dataset_name hf-internal-testing/librispeech_asr_dummy
            --dataset_config_name clean
            --train_split_name validation
            --validation_split_name validation
            --output_dir {output_dir}
            --num_train_epochs {str(UpperCamelCase_				)}
            --per_device_train_batch_size 2
            --per_device_eval_batch_size 2
            --evaluation_strategy steps
            --learning_rate 5e-4
            --warmup_steps 8
            --orthography timit
            --preprocessing_num_workers 1
            --group_by_length
            --freeze_feature_extractor
            --report_to none
            --save_steps 0
            --eval_steps {eval_steps}
            --report_to none
        '''.split()

											if fpaa:
																	args.extend(['''--fp16''']				)

											# currently ds_config_wav2vec2_zero.json requires "zero_optimization.find_unused_parameters": true,
											# hence the separate config files
											lowerCAmelCase					:	Optional[int]       =   F'''--deepspeed {self.test_file_dir_str}/ds_config_wav2vec2_{stage}.json'''.split()
											lowerCAmelCase					:	Dict       =   [F'''{self.examples_dir_str}/research_projects/wav2vec2/run_asr.py''']
											lowerCAmelCase					:	Optional[int]       =   self.get_launcher(UpperCamelCase_				)

											lowerCAmelCase					:	Tuple       =   launcher + script + args + ds_args
											# keep for quick debug
											# print(" ".join([f"\nPYTHONPATH={self.src_dir_str}"] +cmd)); die
											execute_subprocess_async(UpperCamelCase_					,				env=self.get_env()				)

											return output_dir






					def      lowerCamelCase__     (	self  :					List[str]					,				UpperCamelCase_  :					Optional[int]=False				):
											# 1. explicitly set --num_nodes=1 just in case these tests end up run on a multi-node setup
											# - it won't be able to handle that
											# 2. for now testing with just 2 gpus max (since some quality tests may give different
											# results with mode gpus because we use very little data)
											lowerCAmelCase					:	Union[str, Any]       =   min(2					,				get_gpu_count()				) if distributed else 1
											return F'''deepspeed --num_nodes 1 --num_gpus {num_gpus}'''.split()






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


