


"""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  UpperCamelCase_	:

	"""simple docstring"""





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




	def   UpperCAmelCase_      ( self			:				Optional[int]   )       ->	int:
						__SCREAMING_SNAKE_CASE						   =				ids_tensor([self.batch_size, self.seq_length]			,			self.vocab_size   )

						__SCREAMING_SNAKE_CASE						   =				None
						if self.use_input_mask:
											__SCREAMING_SNAKE_CASE						   =				random_attention_mask([self.batch_size, self.seq_length]   )

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

						__SCREAMING_SNAKE_CASE						   =				None
						__SCREAMING_SNAKE_CASE						   =				None
						__SCREAMING_SNAKE_CASE						   =				None
						if self.use_labels:
											__SCREAMING_SNAKE_CASE						   =				ids_tensor([self.batch_size]			,			self.type_sequence_label_size   )
											__SCREAMING_SNAKE_CASE						   =				ids_tensor([self.batch_size, self.seq_length]			,			self.num_labels   )
											__SCREAMING_SNAKE_CASE						   =				ids_tensor([self.batch_size]			,			self.num_choices   )

						__SCREAMING_SNAKE_CASE						   =				self.get_config()

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




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




	def   UpperCAmelCase_      ( self			:				Optional[int]			,			UpperCAmelCase__			:				Dict			,			UpperCAmelCase__			:				List[str]			,			UpperCAmelCase__			:				str			,			UpperCAmelCase__			:				List[Any]			,			UpperCAmelCase__			:				List[Any]			,			UpperCAmelCase__			:				List[str]			,			UpperCAmelCase__			:				Optional[Any]			,			UpperCAmelCase__			:				Any			,			UpperCAmelCase__			:				Dict			,			)       ->	Optional[Any]:
						__SCREAMING_SNAKE_CASE						   =				True
						__SCREAMING_SNAKE_CASE						   =				LlamaModel(UpperCAmelCase__   )
						model.to(UpperCAmelCase__   )
						model.eval()
						__SCREAMING_SNAKE_CASE						   =				model(
						    UpperCAmelCase__			,			attention_mask=UpperCAmelCase__			,			encoder_hidden_states=UpperCAmelCase__			,			encoder_attention_mask=UpperCAmelCase__			,			)
						__SCREAMING_SNAKE_CASE						   =				model(
						    UpperCAmelCase__			,			attention_mask=UpperCAmelCase__			,			encoder_hidden_states=UpperCAmelCase__			,			)
						__SCREAMING_SNAKE_CASE						   =				model(UpperCAmelCase__			,			attention_mask=UpperCAmelCase__   )
						self.parent.assertEqual(result.last_hidden_state.shape			,			(self.batch_size, self.seq_length, self.hidden_size)   )




	def   UpperCAmelCase_      ( self			:				Optional[Any]			,			UpperCAmelCase__			:				int			,			UpperCAmelCase__			:				Any			,			UpperCAmelCase__			:				Optional[int]			,			UpperCAmelCase__			:				Dict			,			UpperCAmelCase__			:				Any			,			UpperCAmelCase__			:				Any			,			UpperCAmelCase__			:				str			,			UpperCAmelCase__			:				Optional[int]			,			UpperCAmelCase__			:				str			,			)       ->	str:
						__SCREAMING_SNAKE_CASE						   =				LlamaForCausalLM(config=UpperCAmelCase__   )
						model.to(UpperCAmelCase__   )
						model.eval()
						__SCREAMING_SNAKE_CASE						   =				model(UpperCAmelCase__			,			attention_mask=UpperCAmelCase__			,			labels=UpperCAmelCase__   )
						self.parent.assertEqual(result.logits.shape			,			(self.batch_size, self.seq_length, self.vocab_size)   )




	def   UpperCAmelCase_      ( self			:				str			,			UpperCAmelCase__			:				Any			,			UpperCAmelCase__			:				Any			,			UpperCAmelCase__			:				List[Any]			,			UpperCAmelCase__			:				Tuple			,			UpperCAmelCase__			:				Any			,			UpperCAmelCase__			:				int			,			UpperCAmelCase__			:				List[str]			,			UpperCAmelCase__			:				Union[str, Any]			,			UpperCAmelCase__			:				Union[str, Any]			,			)       ->	Any:
						__SCREAMING_SNAKE_CASE						   =				True
						__SCREAMING_SNAKE_CASE						   =				True
						__SCREAMING_SNAKE_CASE						   =				LlamaForCausalLM(config=UpperCAmelCase__   )
						model.to(UpperCAmelCase__   )
						model.eval()

						# first forward pass
						__SCREAMING_SNAKE_CASE						   =				model(
						    UpperCAmelCase__			,			attention_mask=UpperCAmelCase__			,			encoder_hidden_states=UpperCAmelCase__			,			encoder_attention_mask=UpperCAmelCase__			,			use_cache=UpperCAmelCase__			,			)
						__SCREAMING_SNAKE_CASE						   =				outputs.past_key_values

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

						# append to next input_ids and
						__SCREAMING_SNAKE_CASE						   =				torch.cat([input_ids, next_tokens]			,			dim=-1   )
						__SCREAMING_SNAKE_CASE						   =				torch.cat([input_mask, next_mask]			,			dim=-1   )

						__SCREAMING_SNAKE_CASE						   =				model(
						    UpperCAmelCase__			,			attention_mask=UpperCAmelCase__			,			encoder_hidden_states=UpperCAmelCase__			,			encoder_attention_mask=UpperCAmelCase__			,			output_hidden_states=UpperCAmelCase__			,			)["hidden_states"][0]
						__SCREAMING_SNAKE_CASE						   =				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
						__SCREAMING_SNAKE_CASE						   =				ids_tensor((1,)			,			output_from_past.shape[-1]   ).item()
						__SCREAMING_SNAKE_CASE						   =				output_from_no_past[:, -3:, random_slice_idx].detach()
						__SCREAMING_SNAKE_CASE						   =				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   UpperCAmelCase_      ( self			:				str   )       ->	Optional[int]:
						__SCREAMING_SNAKE_CASE						   =				self.prepare_config_and_inputs()
						(
						    (
						    __SCREAMING_SNAKE_CASE
						)   ,  (
						    __SCREAMING_SNAKE_CASE
						)   ,  (
						    __SCREAMING_SNAKE_CASE
						)   ,  (
						    __SCREAMING_SNAKE_CASE
						)   ,  (
						    __SCREAMING_SNAKE_CASE
						)   ,  (
						    __SCREAMING_SNAKE_CASE
						)   ,  (
						    __SCREAMING_SNAKE_CASE
						)   ,  
						)						   =				config_and_inputs
						__SCREAMING_SNAKE_CASE						   =				{"input_ids": input_ids, "attention_mask": input_mask}
						return config, inputs_dict




@require_torch
class  UpperCamelCase_	( UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  unittest.TestCase):

	"""simple docstring"""





	snake_case__				:  Dict    					=			(LlamaModel, LlamaForCausalLM, LlamaForSequenceClassification) if is_torch_available() else ()
	snake_case__				:  str    					=			(LlamaForCausalLM,) if is_torch_available() else ()
	snake_case__				:  List[Any]    					=			(
	    {
	        "feature-extraction": LlamaModel,
	        "text-classification": LlamaForSequenceClassification,
	        "text-generation": LlamaForCausalLM,
	        "zero-shot": LlamaForSequenceClassification,
	    }
	    if is_torch_available()
	    else {}
	)
	snake_case__				:  Dict    					=			False
	snake_case__				:  Optional[Any]    					=			False




	def   UpperCAmelCase_      ( self			:				int   )       ->	Tuple:
						__SCREAMING_SNAKE_CASE						   =				LlamaModelTester(self   )
						__SCREAMING_SNAKE_CASE						   =				ConfigTester(self			,			config_class=UpperCAmelCase__			,			hidden_size=3_7   )




	def   UpperCAmelCase_      ( self			:				List[str]   )       ->	Optional[Any]:
						self.config_tester.run_common_tests()




	def   UpperCAmelCase_      ( self			:				str   )       ->	Optional[int]:
						__SCREAMING_SNAKE_CASE						   =				self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*UpperCAmelCase__   )




	def   UpperCAmelCase_      ( self			:				Any   )       ->	Optional[Any]:
						__SCREAMING_SNAKE_CASE						   =				self.model_tester.prepare_config_and_inputs()
						for type in ["absolute", "relative_key", "relative_key_query"]:
											__SCREAMING_SNAKE_CASE						   =				type
											self.model_tester.create_and_check_model(*UpperCAmelCase__   )




	def   UpperCAmelCase_      ( self			:				List[Any]   )       ->	List[str]:
						__SCREAMING_SNAKE_CASE   ,  __SCREAMING_SNAKE_CASE						   =				self.model_tester.prepare_config_and_inputs_for_common()
						__SCREAMING_SNAKE_CASE						   =				3
						__SCREAMING_SNAKE_CASE						   =				input_dict["input_ids"]
						__SCREAMING_SNAKE_CASE						   =				input_ids.ne(1   ).to(UpperCAmelCase__   )
						__SCREAMING_SNAKE_CASE						   =				ids_tensor([self.model_tester.batch_size]			,			self.model_tester.type_sequence_label_size   )
						__SCREAMING_SNAKE_CASE						   =				LlamaForSequenceClassification(UpperCAmelCase__   )
						model.to(UpperCAmelCase__   )
						model.eval()
						__SCREAMING_SNAKE_CASE						   =				model(UpperCAmelCase__			,			attention_mask=UpperCAmelCase__			,			labels=UpperCAmelCase__   )
						self.assertEqual(result.logits.shape			,			(self.model_tester.batch_size, self.model_tester.num_labels)   )




	def   UpperCAmelCase_      ( self			:				Optional[Any]   )       ->	Union[str, Any]:
						__SCREAMING_SNAKE_CASE   ,  __SCREAMING_SNAKE_CASE						   =				self.model_tester.prepare_config_and_inputs_for_common()
						__SCREAMING_SNAKE_CASE						   =				3
						__SCREAMING_SNAKE_CASE						   =				"single_label_classification"
						__SCREAMING_SNAKE_CASE						   =				input_dict["input_ids"]
						__SCREAMING_SNAKE_CASE						   =				input_ids.ne(1   ).to(UpperCAmelCase__   )
						__SCREAMING_SNAKE_CASE						   =				ids_tensor([self.model_tester.batch_size]			,			self.model_tester.type_sequence_label_size   )
						__SCREAMING_SNAKE_CASE						   =				LlamaForSequenceClassification(UpperCAmelCase__   )
						model.to(UpperCAmelCase__   )
						model.eval()
						__SCREAMING_SNAKE_CASE						   =				model(UpperCAmelCase__			,			attention_mask=UpperCAmelCase__			,			labels=UpperCAmelCase__   )
						self.assertEqual(result.logits.shape			,			(self.model_tester.batch_size, self.model_tester.num_labels)   )




	def   UpperCAmelCase_      ( self			:				List[Any]   )       ->	Dict:
						__SCREAMING_SNAKE_CASE   ,  __SCREAMING_SNAKE_CASE						   =				self.model_tester.prepare_config_and_inputs_for_common()
						__SCREAMING_SNAKE_CASE						   =				3
						__SCREAMING_SNAKE_CASE						   =				"multi_label_classification"
						__SCREAMING_SNAKE_CASE						   =				input_dict["input_ids"]
						__SCREAMING_SNAKE_CASE						   =				input_ids.ne(1   ).to(UpperCAmelCase__   )
						__SCREAMING_SNAKE_CASE						   =				ids_tensor(
						    [self.model_tester.batch_size, config.num_labels]			,			self.model_tester.type_sequence_label_size   ).to(torch.float   )
						__SCREAMING_SNAKE_CASE						   =				LlamaForSequenceClassification(UpperCAmelCase__   )
						model.to(UpperCAmelCase__   )
						model.eval()
						__SCREAMING_SNAKE_CASE						   =				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   UpperCAmelCase_      ( self			:				List[Any]   )       ->	List[str]:
						pass




	@parameterized.expand([("linear",), ("dynamic",)]   )
	def   UpperCAmelCase_      ( self			:				int			,			UpperCAmelCase__			:				Optional[Any]   )       ->	Dict:
						__SCREAMING_SNAKE_CASE   ,  __SCREAMING_SNAKE_CASE						   =				self.model_tester.prepare_config_and_inputs_for_common()
						__SCREAMING_SNAKE_CASE						   =				ids_tensor([1, 1_0]			,			config.vocab_size   )
						__SCREAMING_SNAKE_CASE						   =				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
						__SCREAMING_SNAKE_CASE						   =				LlamaModel(UpperCAmelCase__   )
						original_model.to(UpperCAmelCase__   )
						original_model.eval()
						__SCREAMING_SNAKE_CASE						   =				original_model(UpperCAmelCase__   ).last_hidden_state
						__SCREAMING_SNAKE_CASE						   =				original_model(UpperCAmelCase__   ).last_hidden_state

						set_seed(4_2   )  # Fixed seed at init time so the two models get the same random weights
						__SCREAMING_SNAKE_CASE						   =				{"type": scaling_type, "factor": 10.0}
						__SCREAMING_SNAKE_CASE						   =				LlamaModel(UpperCAmelCase__   )
						scaled_model.to(UpperCAmelCase__   )
						scaled_model.eval()
						__SCREAMING_SNAKE_CASE						   =				scaled_model(UpperCAmelCase__   ).last_hidden_state
						__SCREAMING_SNAKE_CASE						   =				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  UpperCamelCase_	( unittest.TestCase):

	"""simple docstring"""





	@unittest.skip("Logits are not exactly the same, once we fix the instabalities somehow, will update!"   )
	@slow
	def   UpperCAmelCase_      ( self			:				Tuple   )       ->	Optional[Any]:
						__SCREAMING_SNAKE_CASE						   =				[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]
						__SCREAMING_SNAKE_CASE						   =				LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf"			,			device_map="auto"   )
						__SCREAMING_SNAKE_CASE						   =				model(torch.tensor([input_ids]   )   )
						# Expected mean on dim = -1
						__SCREAMING_SNAKE_CASE						   =				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
						__SCREAMING_SNAKE_CASE						   =				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   UpperCAmelCase_      ( self			:				List[Any]   )       ->	Any:
						__SCREAMING_SNAKE_CASE						   =				[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]
						__SCREAMING_SNAKE_CASE						   =				LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-13b-hf"			,			device_map="auto"   )
						__SCREAMING_SNAKE_CASE						   =				model(torch.tensor(UpperCAmelCase__   )   )
						# Expected mean on dim = -1
						__SCREAMING_SNAKE_CASE						   =				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
						__SCREAMING_SNAKE_CASE						   =				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   UpperCAmelCase_      ( self			:				Dict   )       ->	int:
						__SCREAMING_SNAKE_CASE						   =				[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]
						__SCREAMING_SNAKE_CASE						   =				LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-13b-chat-hf"			,			device_map="auto"   )
						__SCREAMING_SNAKE_CASE						   =				model(torch.tensor(UpperCAmelCase__   )   )
						# Expected mean on dim = -1
						__SCREAMING_SNAKE_CASE						   =				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
						__SCREAMING_SNAKE_CASE						   =				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   UpperCAmelCase_      ( self			:				Dict   )       ->	int:
						__SCREAMING_SNAKE_CASE						   =				[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]
						__SCREAMING_SNAKE_CASE						   =				LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-70b-hf"			,			device_map="auto"   )
						__SCREAMING_SNAKE_CASE						   =				model(torch.tensor(UpperCAmelCase__   )   )

						__SCREAMING_SNAKE_CASE						   =				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
						__SCREAMING_SNAKE_CASE						   =				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   UpperCAmelCase_      ( self			:				Any   )       ->	Optional[Any]:
						__SCREAMING_SNAKE_CASE						   =				"Simply put, the theory of relativity states that 1) the laws of physics are the same everywhere in the universe and 2) the passage of time and the length of objects can vary depending on the observer\'s frame of reference.\n\nThe first part of the theory, that the laws of physics are the same everywhere, is known as the \"princi"
						__SCREAMING_SNAKE_CASE						   =				"Simply put, the theory of relativity states that "
						__SCREAMING_SNAKE_CASE						   =				LlamaTokenizer.from_pretrained("meta-llama/Llama-2-13b-chat-hf"   )
						__SCREAMING_SNAKE_CASE						   =				tokenizer.encode(UpperCAmelCase__			,			return_tensors="pt"   )
						__SCREAMING_SNAKE_CASE						   =				LlamaForCausalLM.from_pretrained(
						    "meta-llama/Llama-2-13b-chat-hf"			,			device_map="sequential"			,			use_safetensors=UpperCAmelCase__   )

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




"""simple docstring"""







import importlib.util
import json
import os
import warnings
from dataclasses import dataclass, field

import torch

from ..training_args import TrainingArguments
from ..utils import cached_property, is_sagemaker_dp_enabled, logging


a__				:							Union[str, Any]						=     logging.get_logger(__name__)




def    UpperCAmelCase__       ():




					'''simple docstring'''


					__SCREAMING_SNAKE_CASE						   =				os.getenv("SM_HP_MP_PARAMETERS"  ,						"{}"   )
					try:
										# Parse it and check the field "partitions" is included, it is required for model parallel.
										__SCREAMING_SNAKE_CASE						   =				json.loads(lowerCAmelCase_   )
										if "partitions" not in smp_options:
															return False
					except json.JSONDecodeError:
										return False

					# Get the sagemaker specific framework parameters from mpi_options variable.
					__SCREAMING_SNAKE_CASE						   =				os.getenv("SM_FRAMEWORK_PARAMS"  ,						"{}"   )
					try:
										# Parse it and check the field "sagemaker_distributed_dataparallel_enabled".
										__SCREAMING_SNAKE_CASE						   =				json.loads(lowerCAmelCase_   )
										if not mpi_options.get("sagemaker_mpi_enabled"  ,						lowerCAmelCase_   ):
															return False
					except json.JSONDecodeError:
										return False
					# Lastly, check if the `smdistributed` module is present.
					return importlib.util.find_spec("smdistributed"   ) is not None


if is_sagemaker_model_parallel_available():
		import smdistributed.modelparallel.torch as smp

		smp.init()




@dataclass
class  UpperCamelCase_	( UpperCamelCase):

	"""simple docstring"""





	snake_case__				:  str    					=			field(
	    default=""			,  metadata={"help": "Used by the SageMaker launcher to send mp-specific args. Ignored in SageMakerTrainer"}			,  )




	def   UpperCAmelCase_      ( self			:				List[str]   )       ->	Any:
						super().__post_init__()
						warnings.warn(
						    "`SageMakerTrainingArguments` is deprecated and will be removed in v5 of Transformers. You can use "
						    "`TrainingArguments` instead."			,			UpperCAmelCase__			,			)




	@cached_property
	def   UpperCAmelCase_      ( self			:				List[str]   )       ->	"torch.device":
						logger.info("PyTorch: setting up devices"   )
						if torch.distributed.is_available() and torch.distributed.is_initialized() and self.local_rank == -1:
											logger.warning(
											    "torch.distributed process group is initialized, but local_rank == -1. "
											    "In order to use Torch DDP, launch your script with `python -m torch.distributed.launch"   )
						if self.no_cuda:
											__SCREAMING_SNAKE_CASE						   =				torch.device("cpu"   )
											__SCREAMING_SNAKE_CASE						   =				0
						elif is_sagemaker_model_parallel_available():
											__SCREAMING_SNAKE_CASE						   =				smp.local_rank()
											__SCREAMING_SNAKE_CASE						   =				torch.device("cuda"			,			UpperCAmelCase__   )
											__SCREAMING_SNAKE_CASE						   =				1
						elif is_sagemaker_dp_enabled():
											import smdistributed.dataparallel.torch.torch_smddp  # noqa: F401

											torch.distributed.init_process_group(backend="smddp"			,			timeout=self.ddp_timeout_delta   )
											__SCREAMING_SNAKE_CASE						   =				int(os.getenv("SMDATAPARALLEL_LOCAL_RANK"   )   )
											__SCREAMING_SNAKE_CASE						   =				torch.device("cuda"			,			self.local_rank   )
											__SCREAMING_SNAKE_CASE						   =				1
						elif self.local_rank == -1:
											# if n_gpu is > 1 we'll use nn.DataParallel.
											# If you only want to use a specific subset of GPUs use `CUDA_VISIBLE_DEVICES=0`
											# Explicitly set CUDA to the first (index 0) CUDA device, otherwise `set_device` will
											# trigger an error that a device index is missing. Index 0 takes into account the
											# GPUs available in the environment, so `CUDA_VISIBLE_DEVICES=1,2` with `cuda:0`
											# will use the first GPU in that env, i.e. GPU#1
											__SCREAMING_SNAKE_CASE						   =				torch.device("cuda:0" if torch.cuda.is_available() else "cpu"   )
											# Sometimes the line in the postinit has not been run before we end up here, so just checking we're not at
											# the default value.
											__SCREAMING_SNAKE_CASE						   =				torch.cuda.device_count()
						else:
											# Here, we'll use torch.distributed.
											# Initializes the distributed backend which will take care of synchronizing nodes/GPUs
											if not torch.distributed.is_initialized():
																torch.distributed.init_process_group(backend="nccl"			,			timeout=self.ddp_timeout_delta   )
											__SCREAMING_SNAKE_CASE						   =				torch.device("cuda"			,			self.local_rank   )
											__SCREAMING_SNAKE_CASE						   =				1

						if device.type == "cuda":
											torch.cuda.set_device(UpperCAmelCase__   )

						return device




	@property
	def   UpperCAmelCase_      ( self			:				Dict   )       ->	Any:
						if is_sagemaker_model_parallel_available():
											return smp.dp_size()

						return super().world_size




	@property
	def   UpperCAmelCase_      ( self			:				Union[str, Any]   )       ->	List[Any]:
						return not is_sagemaker_model_parallel_available()




	@property
	def   UpperCAmelCase_      ( self			:				Tuple   )       ->	int:
						return False
