


"""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    :       Tuple    ,    UpperCAmelCase__    :       Any    ,    UpperCAmelCase__    :       Optional[int]=1_3    ,    UpperCAmelCase__    :       str=7    ,    UpperCAmelCase__    :       Optional[Any]=True    ,    UpperCAmelCase__    :       Optional[int]=True    ,    UpperCAmelCase__    :       List[str]=False    ,    UpperCAmelCase__    :       Any=True    ,    UpperCAmelCase__    :       List[Any]=9_9    ,    UpperCAmelCase__    :       Any=3_2    ,    UpperCAmelCase__    :       Optional[Any]=5    ,    UpperCAmelCase__    :       List[str]=4    ,    UpperCAmelCase__    :       Any=3_7    ,    UpperCAmelCase__    :       Any="gelu"    ,    UpperCAmelCase__    :       str=0.1    ,    UpperCAmelCase__    :       Tuple=0.1    ,    UpperCAmelCase__    :       Union[str, Any]=5_1_2    ,    UpperCAmelCase__    :       str=1_6    ,    UpperCAmelCase__    :       str=2    ,    UpperCAmelCase__    :       List[Any]=0.02    ,    UpperCAmelCase__    :       Dict=3    ,    UpperCAmelCase__    :       List[str]=4    ,    UpperCAmelCase__    :       str=None    ,    ) ->			Union[str, Any]:
						__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    :       List[Any]     ) ->			Union[str, Any]:
						__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    :       Any     ) ->			Optional[int]:
						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    :       Tuple    ,    UpperCAmelCase__    :       List[str]    ,    UpperCAmelCase__    :       Dict    ,    UpperCAmelCase__    :       Any    ,    UpperCAmelCase__    :       Union[str, Any]    ,    UpperCAmelCase__    :       List[Any]    ,    UpperCAmelCase__    :       List[Any]    ,    UpperCAmelCase__    :       List[str]     ) ->			str:
						__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    :       Any    ,    UpperCAmelCase__    :       int    ,    UpperCAmelCase__    :       Tuple    ,    UpperCAmelCase__    :       List[str]    ,    UpperCAmelCase__    :       int    ,    UpperCAmelCase__    :       Dict    ,    UpperCAmelCase__    :       Dict    ,    UpperCAmelCase__    :       List[str]    ,    UpperCAmelCase__    :       Any    ,    UpperCAmelCase__    :       List[str]    ,    ) ->			Dict:
						__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[int]    ,    UpperCAmelCase__    :       str    ,    UpperCAmelCase__    :       Optional[int]    ,    UpperCAmelCase__    :       Union[str, Any]    ,    UpperCAmelCase__    :       Optional[int]    ,    UpperCAmelCase__    :       Optional[Any]    ,    UpperCAmelCase__    :       int    ,    UpperCAmelCase__    :       Optional[Any]    ,    UpperCAmelCase__    :       List[str]    ,    UpperCAmelCase__    :       int    ,    ) ->			Optional[Any]:
						__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    :       Union[str, Any]    ,    UpperCAmelCase__    :       List[str]    ,    UpperCAmelCase__    :       List[str]    ,    UpperCAmelCase__    :       List[Any]    ,    UpperCAmelCase__    :       Optional[Any]    ,    UpperCAmelCase__    :       Tuple    ,    UpperCAmelCase__    :       Tuple    ,    UpperCAmelCase__    :       Optional[int]    ,    UpperCAmelCase__    :       Dict    ,    UpperCAmelCase__    :       List[Any]    ,    ) ->			List[str]:
						__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    :       Any     ) ->			str:
						__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__       :						Optional[Any]					      =    (LlamaModel, LlamaForCausalLM, LlamaForSequenceClassification) if is_torch_available() else ()
	snake_case__       :						Optional[Any]					      =    (LlamaForCausalLM,) if is_torch_available() else ()
	snake_case__       :						int					      =    (
	    {
	        "feature-extraction": LlamaModel,
	        "text-classification": LlamaForSequenceClassification,
	        "text-generation": LlamaForCausalLM,
	        "zero-shot": LlamaForSequenceClassification,
	    }
	    if is_torch_available()
	    else {}
	)
	snake_case__       :						Any					      =    False
	snake_case__       :						int					      =    False




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




	def 						UpperCAmelCase_   (				self    :       Any     ) ->			Dict:
						self.config_tester.run_common_tests()




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




	def 						UpperCAmelCase_   (				self    :       Tuple     ) ->			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    :       Optional[Any]     ) ->			int:
						__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    :       Union[str, Any]     ) ->			Tuple:
						__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    :       str     ) ->			int:
						__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    :       str     ) ->			List[str]:
						pass




	@parameterized.expand([("linear",), ("dynamic",)]     )
	def 						UpperCAmelCase_   (				self    :       int    ,    UpperCAmelCase__    :       List[Any]     ) ->			str:
						__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    :       Dict     ) ->			Tuple:
						__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    :       Optional[int]     ) ->			List[str]:
						__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    :       List[Any]     ) ->			Dict:
						__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    :       int     ) ->			str:
						__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    :       str     ) ->			Optional[int]:
						__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"""







from __future__ import annotations

import math




class 				UpperCamelCase_			:

	"""simple docstring"""





	def __init__(				self    :       Union[str, Any]    ,    UpperCAmelCase__    :       int     ) ->			None:
						__SCREAMING_SNAKE_CASE		     =       size
						# approximate the overall size of segment tree with given value
						__SCREAMING_SNAKE_CASE		     =       [0 for i in range(0    ,    4 * size     )]
						# create array to store lazy update
						__SCREAMING_SNAKE_CASE		     =       [0 for i in range(0    ,    4 * size     )]
						__SCREAMING_SNAKE_CASE		     =       [0 for i in range(0    ,    4 * size     )]  # flag for lazy update




	def 						UpperCAmelCase_   (				self    :       str    ,    UpperCAmelCase__    :       int     ) ->			int:
						return idx * 2




	def 						UpperCAmelCase_   (				self    :       Any    ,    UpperCAmelCase__    :       int     ) ->			int:
						return idx * 2 + 1




	def 						UpperCAmelCase_   (				self    :       List[Any]    ,    UpperCAmelCase__    :       int    ,    UpperCAmelCase__    :       int    ,    UpperCAmelCase__    :       int    ,    UpperCAmelCase__    :       list[int]     ) ->			None:
						if left_element == right_element:
											__SCREAMING_SNAKE_CASE		     =       a[left_element - 1]
						else:
											__SCREAMING_SNAKE_CASE		     =       (left_element + right_element) // 2
											self.build(self.left(UpperCAmelCase__     )    ,    UpperCAmelCase__    ,    UpperCAmelCase__    ,    UpperCAmelCase__     )
											self.build(self.right(UpperCAmelCase__     )    ,    mid + 1    ,    UpperCAmelCase__    ,    UpperCAmelCase__     )
											__SCREAMING_SNAKE_CASE		     =       max(
											    self.segment_tree[self.left(UpperCAmelCase__     )]    ,    self.segment_tree[self.right(UpperCAmelCase__     )]     )




	def 						UpperCAmelCase_   (				self    :       Tuple    ,    UpperCAmelCase__    :       int    ,    UpperCAmelCase__    :       int    ,    UpperCAmelCase__    :       int    ,    UpperCAmelCase__    :       int    ,    UpperCAmelCase__    :       int    ,    UpperCAmelCase__    :       int     ) ->			bool:
						if self.flag[idx] is True:
											__SCREAMING_SNAKE_CASE		     =       self.lazy[idx]
											__SCREAMING_SNAKE_CASE		     =       False
											if left_element != right_element:
																__SCREAMING_SNAKE_CASE		     =       self.lazy[idx]
																__SCREAMING_SNAKE_CASE		     =       self.lazy[idx]
																__SCREAMING_SNAKE_CASE		     =       True
																__SCREAMING_SNAKE_CASE		     =       True

						if right_element < a or left_element > b:
											return True
						if left_element >= a and right_element <= b:
											__SCREAMING_SNAKE_CASE		     =       val
											if left_element != right_element:
																__SCREAMING_SNAKE_CASE		     =       val
																__SCREAMING_SNAKE_CASE		     =       val
																__SCREAMING_SNAKE_CASE		     =       True
																__SCREAMING_SNAKE_CASE		     =       True
											return True
						__SCREAMING_SNAKE_CASE		     =       (left_element + right_element) // 2
						self.update(self.left(UpperCAmelCase__     )    ,    UpperCAmelCase__    ,    UpperCAmelCase__    ,    UpperCAmelCase__    ,    UpperCAmelCase__    ,    UpperCAmelCase__     )
						self.update(self.right(UpperCAmelCase__     )    ,    mid + 1    ,    UpperCAmelCase__    ,    UpperCAmelCase__    ,    UpperCAmelCase__    ,    UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       max(
						    self.segment_tree[self.left(UpperCAmelCase__     )]    ,    self.segment_tree[self.right(UpperCAmelCase__     )]     )
						return True




	def 						UpperCAmelCase_   (				self    :       int    ,    UpperCAmelCase__    :       int    ,    UpperCAmelCase__    :       int    ,    UpperCAmelCase__    :       int    ,    UpperCAmelCase__    :       int    ,    UpperCAmelCase__    :       int     ) ->			int | float:
						if self.flag[idx] is True:
											__SCREAMING_SNAKE_CASE		     =       self.lazy[idx]
											__SCREAMING_SNAKE_CASE		     =       False
											if left_element != right_element:
																__SCREAMING_SNAKE_CASE		     =       self.lazy[idx]
																__SCREAMING_SNAKE_CASE		     =       self.lazy[idx]
																__SCREAMING_SNAKE_CASE		     =       True
																__SCREAMING_SNAKE_CASE		     =       True
						if right_element < a or left_element > b:
											return -math.inf
						if left_element >= a and right_element <= b:
											return self.segment_tree[idx]
						__SCREAMING_SNAKE_CASE		     =       (left_element + right_element) // 2
						__SCREAMING_SNAKE_CASE		     =       self.query(self.left(UpperCAmelCase__     )    ,    UpperCAmelCase__    ,    UpperCAmelCase__    ,    UpperCAmelCase__    ,    UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       self.query(self.right(UpperCAmelCase__     )    ,    mid + 1    ,    UpperCAmelCase__    ,    UpperCAmelCase__    ,    UpperCAmelCase__     )
						return max(UpperCAmelCase__    ,    UpperCAmelCase__     )




	def __str__(				self    :       int     ) ->			str:
						return str([self.query(1    ,    1    ,    self.size    ,    UpperCAmelCase__    ,    UpperCAmelCase__     ) for i in range(1    ,    self.size + 1     )]     )


if __name__ == "__main__":
		a__       :							Tuple        =  [1, 2, -4, 7, 3, -5, 6, 1_1, -2_0, 9, 1_4, 1_5, 5, 2, -8]
		a__       :							Dict        =  1_5
		a__       :							int        =  SegmentTree(size)
		segt.build(1, 1, size, A)
		print(segt.query(1, 1, size, 4, 6))
		print(segt.query(1, 1, size, 7, 1_1))
		print(segt.query(1, 1, size, 7, 1_2))
		segt.update(1, 1, size, 1, 3, 1_1_1)
		print(segt.query(1, 1, size, 1, 1_5))
		segt.update(1, 1, size, 7, 8, 2_3_5)
		print(segt)
