

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 					__lowercase  :





							"""simple docstring"""

							def __init__(					self    ,						__UpperCAmelCase    ,						__UpperCAmelCase=13    ,						__UpperCAmelCase=7    ,						__UpperCAmelCase=True    ,						__UpperCAmelCase=True    ,						__UpperCAmelCase=False    ,						__UpperCAmelCase=True    ,						__UpperCAmelCase=99    ,						__UpperCAmelCase=32    ,						__UpperCAmelCase=5    ,						__UpperCAmelCase=4    ,						__UpperCAmelCase=37    ,						__UpperCAmelCase="gelu"    ,						__UpperCAmelCase=0.1    ,						__UpperCAmelCase=0.1    ,						__UpperCAmelCase=5_12    ,						__UpperCAmelCase=16    ,						__UpperCAmelCase=2    ,						__UpperCAmelCase=0.0_2    ,						__UpperCAmelCase=3    ,						__UpperCAmelCase=4    ,						__UpperCAmelCase=None    ,						)		->       str:
														A      :					List[str]      		=					parent
														A      :					List[Any]      		=					batch_size
														A      :					int      		=					seq_length
														A      :					Union[str, Any]      		=					is_training
														A      :					Tuple      		=					use_input_mask
														A      :					List[str]      		=					use_token_type_ids
														A      :					Tuple      		=					use_labels
														A      :					Optional[int]      		=					vocab_size
														A      :					Any      		=					hidden_size
														A      :					List[Any]      		=					num_hidden_layers
														A      :					Optional[int]      		=					num_attention_heads
														A      :					Tuple      		=					intermediate_size
														A      :					Dict      		=					hidden_act
														A      :					Any      		=					hidden_dropout_prob
														A      :					int      		=					attention_probs_dropout_prob
														A      :					int      		=					max_position_embeddings
														A      :					Dict      		=					type_vocab_size
														A      :					Dict      		=					type_sequence_label_size
														A      :					Union[str, Any]      		=					initializer_range
														A      :					Tuple      		=					num_labels
														A      :					Any      		=					num_choices
														A      :					str      		=					scope




							def       snake_case			(					self						)		->       Any:
														A      :					Tuple      		=					ids_tensor([self.batch_size, self.seq_length]    ,						self.vocab_size						)

														A      :					Optional[Any]      		=					None
														if self.use_input_mask:
																					A      :					int      		=					random_attention_mask([self.batch_size, self.seq_length]						)

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

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

														A      :					int      		=					self.get_config()

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




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




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




							def       snake_case			(					self    ,						__UpperCAmelCase    ,						__UpperCAmelCase    ,						__UpperCAmelCase    ,						__UpperCAmelCase    ,						__UpperCAmelCase    ,						__UpperCAmelCase    ,						__UpperCAmelCase    ,						__UpperCAmelCase    ,						__UpperCAmelCase    ,						)		->       int:
														A      :					Union[str, Any]      		=					LlamaForCausalLM(config=__UpperCAmelCase						)
														model.to(__UpperCAmelCase						)
														model.eval()
														A      :					str      		=					model(__UpperCAmelCase    ,						attention_mask=__UpperCAmelCase    ,						labels=__UpperCAmelCase						)
														self.parent.assertEqual(result.logits.shape    ,						(self.batch_size, self.seq_length, self.vocab_size)						)




							def       snake_case			(					self    ,						__UpperCAmelCase    ,						__UpperCAmelCase    ,						__UpperCAmelCase    ,						__UpperCAmelCase    ,						__UpperCAmelCase    ,						__UpperCAmelCase    ,						__UpperCAmelCase    ,						__UpperCAmelCase    ,						__UpperCAmelCase    ,						)		->       Optional[int]:
														A      :					Union[str, Any]      		=					True
														A      :					Any      		=					True
														A      :					Optional[int]      		=					LlamaForCausalLM(config=__UpperCAmelCase						)
														model.to(__UpperCAmelCase						)
														model.eval()

														# first forward pass
														A      :					List[str]      		=					model(
														    __UpperCAmelCase    ,						attention_mask=__UpperCAmelCase    ,						encoder_hidden_states=__UpperCAmelCase    ,						encoder_attention_mask=__UpperCAmelCase    ,						use_cache=__UpperCAmelCase    ,						)
														A      :					Tuple      		=					outputs.past_key_values

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

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

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







@require_torch
class 					__lowercase  (						_SCREAMING_SNAKE_CASE		,       _SCREAMING_SNAKE_CASE		,       _SCREAMING_SNAKE_CASE		,       unittest.TestCase				):





							"""simple docstring"""

							UpperCAmelCase_						:		Tuple								=	(LlamaModel, LlamaForCausalLM, LlamaForSequenceClassification) if is_torch_available() else ()
							UpperCAmelCase_						:		int								=	(LlamaForCausalLM,) if is_torch_available() else ()
							UpperCAmelCase_						:		Optional[int]								=	(
							    {
							        '''feature-extraction''': LlamaModel,
							        '''text-classification''': LlamaForSequenceClassification,
							        '''text-generation''': LlamaForCausalLM,
							        '''zero-shot''': LlamaForSequenceClassification,
							    }
							    if is_torch_available()
							    else {}
							)
							UpperCAmelCase_						:		Union[str, Any]								=	False
							UpperCAmelCase_						:		Optional[Any]								=	False




							def       snake_case			(					self						)		->       int:
														A      :					Dict      		=					LlamaModelTester(self						)
														A      :					Tuple      		=					ConfigTester(self    ,						config_class=__UpperCAmelCase    ,						hidden_size=37						)




							def       snake_case			(					self						)		->       Dict:
														self.config_tester.run_common_tests()




							def       snake_case			(					self						)		->       Optional[int]:
														A      :					Any      		=					self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_model(*__UpperCAmelCase						)




							def       snake_case			(					self						)		->       Optional[int]:
														A      :					List[Any]      		=					self.model_tester.prepare_config_and_inputs()
														for type in ["absolute", "relative_key", "relative_key_query"]:
																					A      :					str      		=					type
																					self.model_tester.create_and_check_model(*__UpperCAmelCase						)




							def       snake_case			(					self						)		->       Optional[int]:
														A		,		A      :					Optional[Any]      		=					self.model_tester.prepare_config_and_inputs_for_common()
														A      :					List[Any]      		=					3
														A      :					str      		=					input_dict['''input_ids''']
														A      :					str      		=					input_ids.ne(1						).to(__UpperCAmelCase						)
														A      :					Dict      		=					ids_tensor([self.model_tester.batch_size]    ,						self.model_tester.type_sequence_label_size						)
														A      :					Dict      		=					LlamaForSequenceClassification(__UpperCAmelCase						)
														model.to(__UpperCAmelCase						)
														model.eval()
														A      :					Dict      		=					model(__UpperCAmelCase    ,						attention_mask=__UpperCAmelCase    ,						labels=__UpperCAmelCase						)
														self.assertEqual(result.logits.shape    ,						(self.model_tester.batch_size, self.model_tester.num_labels)						)




							def       snake_case			(					self						)		->       int:
														A		,		A      :					Optional[Any]      		=					self.model_tester.prepare_config_and_inputs_for_common()
														A      :					Union[str, Any]      		=					3
														A      :					Union[str, Any]      		=					'''single_label_classification'''
														A      :					List[Any]      		=					input_dict['''input_ids''']
														A      :					Any      		=					input_ids.ne(1						).to(__UpperCAmelCase						)
														A      :					Tuple      		=					ids_tensor([self.model_tester.batch_size]    ,						self.model_tester.type_sequence_label_size						)
														A      :					Any      		=					LlamaForSequenceClassification(__UpperCAmelCase						)
														model.to(__UpperCAmelCase						)
														model.eval()
														A      :					Optional[int]      		=					model(__UpperCAmelCase    ,						attention_mask=__UpperCAmelCase    ,						labels=__UpperCAmelCase						)
														self.assertEqual(result.logits.shape    ,						(self.model_tester.batch_size, self.model_tester.num_labels)						)




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






							@parameterized.expand([('''linear''',), ('''dynamic''',)]						)
							def       snake_case			(					self    ,						__UpperCAmelCase						)		->       Optional[int]:
														A		,		A      :					List[Any]      		=					self.model_tester.prepare_config_and_inputs_for_common()
														A      :					Dict      		=					ids_tensor([1, 10]    ,						config.vocab_size						)
														A      :					Dict      		=					ids_tensor([1, int(config.max_position_embeddings * 1.5						)]    ,						config.vocab_size						)

														set_seed(42						)  # Fixed seed at init time so the two models get the same random weights
														A      :					str      		=					LlamaModel(__UpperCAmelCase						)
														original_model.to(__UpperCAmelCase						)
														original_model.eval()
														A      :					Union[str, Any]      		=					original_model(__UpperCAmelCase						).last_hidden_state
														A      :					Any      		=					original_model(__UpperCAmelCase						).last_hidden_state

														set_seed(42						)  # Fixed seed at init time so the two models get the same random weights
														A      :					Dict      		=					{'''type''': scaling_type, '''factor''': 1_0.0}
														A      :					Union[str, Any]      		=					LlamaModel(__UpperCAmelCase						)
														scaled_model.to(__UpperCAmelCase						)
														scaled_model.eval()
														A      :					Optional[int]      		=					scaled_model(__UpperCAmelCase						).last_hidden_state
														A      :					Tuple      		=					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 					__lowercase  (						unittest.TestCase				):





							"""simple docstring"""

							@unittest.skip('''Logits are not exactly the same, once we fix the instabalities somehow, will update!'''						)
							@slow
							def       snake_case			(					self						)		->       Union[str, Any]:
														A      :					Tuple      		=					[1, 3_06, 46_58, 2_78, 65_93, 3_10, 28_34, 3_38]
														A      :					Union[str, Any]      		=					LlamaForCausalLM.from_pretrained('''meta-llama/Llama-2-7b-hf'''    ,						device_map='''auto'''						)
														A      :					Dict      		=					model(torch.tensor([input_ids]						)						)
														# Expected mean on dim = -1
														A      :					Optional[Any]      		=					torch.tensor([[-6.6_5_5_0, -4.1_2_2_7, -4.9_8_5_9, -3.2_4_0_6, 0.8_2_6_2, -3.0_0_3_3, 1.2_9_6_4, -3.3_6_9_9]]						)
														torch.testing.assert_close(out.mean(-1						)    ,						__UpperCAmelCase    ,						atol=1E-2    ,						rtol=1E-2						)
														# slicing logits[0, 0, 0:30]
														# fmt: off
														A      :					str      		=					torch.tensor([-1_2.8_2_8_1, -7.4_4_5_3, -0.4_6_3_9, -8.0_6_2_5, -7.2_5_0_0, -8.0_0_0_0, -6.4_8_8_3, -7.7_6_9_5, -7.8_4_3_8, -7.0_3_1_2, -6.2_1_8_8, -7.1_3_2_8, -1.8_4_9_6, 1.9_9_6_1, -8.6_2_5_0, -6.7_2_2_7, -1_2.8_2_8_1, -6.9_4_9_2, -7.0_7_4_2, -7.7_8_5_2, -7.5_8_2_0, -7.9_0_6_2, -6.9_3_7_5, -7.9_8_0_5, -8.3_4_3_8, -8.1_5_6_2, -8.0_4_6_9, -7.6_2_5_0, -7.7_4_2_2, -7.3_3_9_8,]						)
														# fmt: on
														torch.testing.assert_close(out[0, 0, :30]    ,						__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       snake_case			(					self						)		->       int:
														A      :					Union[str, Any]      		=					[1, 3_06, 46_58, 2_78, 65_93, 3_10, 28_34, 3_38]
														A      :					int      		=					LlamaForCausalLM.from_pretrained('''meta-llama/Llama-2-13b-hf'''    ,						device_map='''auto'''						)
														A      :					Tuple      		=					model(torch.tensor(__UpperCAmelCase						)						)
														# Expected mean on dim = -1
														A      :					str      		=					torch.tensor([[-2.0_6_2_2, -1.2_7_9_4, -1.1_6_3_8, -0.9_7_8_8, -1.4_6_0_3, -1.0_2_3_8, -1.7_8_9_3, -1.4_4_1_1]]						)
														torch.testing.assert_close(out.mean(-1						)    ,						__UpperCAmelCase    ,						atol=1E-2    ,						rtol=1E-2						)
														# slicing logits[0, 0, 0:30]
														# fmt: off
														A      :					List[str]      		=					torch.tensor([-8.1_4_0_6, -8.0_5_4_7, 2.7_4_6_1, -1.2_3_4_4, -0.1_4_4_8, -1.8_2_6_2, -1.0_0_2_0, -1.8_1_5_4, -1.6_8_9_5, -1.8_5_1_6, -2.3_5_7_4, -0.9_2_7_7, 3.7_5_9_8, 6.5_7_4_2, -1.2_9_9_8, -0.1_1_7_7, -8.1_4_0_6, -2.9_6_8_8, -2.9_1_9_9, -3.1_6_9_9, -3.5_2_5_4, -2.3_5_5_5, -2.7_9_8_8, -3.4_1_4_1, -2.8_2_6_2, -4.5_1_9_5, -3.3_3_7_9, -3.3_1_6_4, -2.7_8_3_2, -3.0_2_7_3]						)
														# fmt: on
														torch.testing.assert_close(out[0, 0, :30]    ,						__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       snake_case			(					self						)		->       Optional[int]:
														A      :					Optional[int]      		=					[1, 3_06, 46_58, 2_78, 65_93, 3_10, 28_34, 3_38]
														A      :					int      		=					LlamaForCausalLM.from_pretrained('''meta-llama/Llama-2-13b-chat-hf'''    ,						device_map='''auto'''						)
														A      :					Optional[int]      		=					model(torch.tensor(__UpperCAmelCase						)						)
														# Expected mean on dim = -1
														A      :					Union[str, Any]      		=					torch.tensor([[-0.8_5_6_2, -1.8_5_2_0, -0.7_5_5_1, -0.4_1_6_2, -1.5_1_6_1, -1.2_0_3_8, -2.4_8_2_3, -2.3_2_5_4]]						)
														torch.testing.assert_close(out.mean(-1						)    ,						__UpperCAmelCase    ,						atol=1E-2    ,						rtol=1E-2						)
														# slicing logits[0, 0, 0:30]
														# fmt: off
														A      :					Optional[Any]      		=					torch.tensor([-2.2_2_2_7, 4.8_8_2_8, 0.9_0_2_3, -0.4_5_7_8, -0.7_8_7_1, -0.1_0_3_3, -0.6_2_2_1, -0.5_7_8_6, -0.7_8_0_3, -1.0_6_7_4, -1.2_9_2_0, -0.1_5_7_0, 0.8_0_0_8, 2.0_7_2_3, -0.9_4_9_7, 0.2_7_7_1, -2.2_2_2_7, -0.7_6_1_2, -1.4_3_4_6, -1.2_0_6_1, -1.6_4_2_6, -0.3_0_0_0, -0.7_1_3_9, -1.1_9_3_4, -1.8_6_9_1, -1.6_9_7_3, -1.5_9_4_7, -1.2_7_0_5, -0.3_5_2_3, -0.5_5_1_3]						)
														# 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       snake_case			(					self						)		->       Dict:
														A      :					Any      		=					[1, 3_06, 46_58, 2_78, 65_93, 3_10, 28_34, 3_38]
														A      :					Dict      		=					LlamaForCausalLM.from_pretrained('''meta-llama/Llama-2-70b-hf'''    ,						device_map='''auto'''						)
														A      :					Any      		=					model(torch.tensor(__UpperCAmelCase						)						)

														A      :					Any      		=					torch.tensor(
														    [[-4.2_3_2_7, -3.3_3_6_0, -4.6_6_6_5, -4.7_6_3_1, -1.8_1_8_0, -3.4_1_7_0, -1.4_2_1_1, -3.1_8_1_0]]    ,						dtype=torch.floataa						)
														torch.testing.assert_close(out.mean(-1						)    ,						__UpperCAmelCase    ,						atol=1E-2    ,						rtol=1E-2						)
														# fmt: off
														A      :					List[str]      		=					torch.tensor([-9.4_9_2_2, -3.9_5_5_1, 1.7_9_9_8, -5.6_7_5_8, -5.1_0_5_5, -5.8_9_8_4, -4.8_3_2_0, -6.8_0_8_6, -6.5_3_9_1, -5.6_1_7_2, -5.5_8_2_0, -5.5_3_5_2, 1.7_8_8_1, 3.6_2_8_9, -6.5_1_1_7, -3.4_7_8_5, -9.5_0_0_0, -6.0_3_5_2, -6.8_1_2_5, -6.0_1_9_5, -6.6_8_3_6, -5.4_7_2_7, -6.2_8_1_2, -6.0_3_9_1, -7.3_3_9_8, -7.4_2_9_7, -7.4_8_4_4, -6.5_8_2_0, -5.8_7_8_9, -5.5_3_1_2]						)
														# fmt: on
														torch.testing.assert_close(out[0, 0, :30]    ,						__UpperCAmelCase    ,						atol=1E-5    ,						rtol=1E-5						)






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

														# greedy generation outputs
														A      :					List[str]      		=					model.generate(__UpperCAmelCase    ,						max_new_tokens=64    ,						top_p=__UpperCAmelCase    ,						temperature=1    ,						do_sample=__UpperCAmelCase						)
														A      :					Optional[Any]      		=					tokenizer.decode(generated_ids[0]    ,						skip_special_tokens=__UpperCAmelCase						)
														self.assertEqual(__UpperCAmelCase    ,						__UpperCAmelCase						)









import os
import posixpath
import uuid
from dataclasses import dataclass
from typing import TYPE_CHECKING, Iterable, List, Optional, Tuple, Union

import numpy as np
import pyarrow as pa

import datasets
from datasets.arrow_writer import ArrowWriter, ParquetWriter
from datasets.config import MAX_SHARD_SIZE
from datasets.filesystems import (
    is_remote_filesystem,
    rename,
)
from datasets.iterable_dataset import _BaseExamplesIterable
from datasets.utils.py_utils import convert_file_size_to_int


lowercase			:				Dict   =     datasets.utils.logging.get_logger(__name__)

if TYPE_CHECKING:
		import pyspark







@dataclass
class 					__lowercase  (						datasets.BuilderConfig				):





							"""simple docstring"""

							UpperCAmelCase_						:		Optional[datasets.Features]								=	None






def        snake_case__		(    lowerCamelCase_     ,				lowerCamelCase_     ,				):
							import pyspark

							def generate_fn():
														A      :					Dict      		=					df.select('''*'''     ,				pyspark.sql.functions.spark_partition_id().alias('''part_id'''		)		)
														for partition_id in partition_order:
																					A      :					Optional[Any]      		=					df_with_partition_id.select('''*'''		).where(F'part_id = {partition_id}'		).drop('''part_id'''		)
																					A      :					Dict      		=					partition_df.collect()
																					A      :					List[str]      		=					0
																					for row in rows:
																												yield F'{partition_id}_{row_id}', row.asDict()
																												row_id += 1

							return generate_fn







class 					__lowercase  (						_BaseExamplesIterable				):





							"""simple docstring"""

							def __init__(					self    ,						__UpperCAmelCase    ,						__UpperCAmelCase=None    ,						)		->       List[str]:
														A      :					List[Any]      		=					df
														A      :					List[str]      		=					partition_order or range(self.df.rdd.getNumPartitions()						)
														A      :					Optional[int]      		=					_generate_iterable_examples(self.df    ,						self.partition_order						)




							def __iter__(					self						)		->       str:
														yield from self.generate_examples_fn()




							def       snake_case			(					self    ,						__UpperCAmelCase						)		->       "SparkExamplesIterable":
														A      :					Tuple      		=					list(range(self.df.rdd.getNumPartitions()						)						)
														generator.shuffle(__UpperCAmelCase						)
														return SparkExamplesIterable(self.df    ,						partition_order=__UpperCAmelCase						)




							def       snake_case			(					self    ,						__UpperCAmelCase    ,						__UpperCAmelCase						)		->       "SparkExamplesIterable":
														A      :					Tuple      		=					self.split_shard_indices_by_worker(__UpperCAmelCase    ,						__UpperCAmelCase						)
														return SparkExamplesIterable(self.df    ,						partition_order=__UpperCAmelCase						)






							@property
							def       snake_case			(					self						)		->       int:
														return len(self.partition_order						)







class 					__lowercase  (						datasets.DatasetBuilder				):





							"""simple docstring"""

							UpperCAmelCase_						:		List[str]								=	SparkConfig




							def __init__(					self    ,						__UpperCAmelCase    ,						__UpperCAmelCase = None    ,						__UpperCAmelCase = None    ,						**__UpperCAmelCase    ,						)		->       Optional[Any]:
														import pyspark

														A      :					Tuple      		=					pyspark.sql.SparkSession.builder.getOrCreate()
														A      :					List[Any]      		=					df
														A      :					List[Any]      		=					working_dir

														super().__init__(
														    cache_dir=__UpperCAmelCase    ,						config_name=str(self.df.semanticHash()						)    ,						**__UpperCAmelCase    ,						)




							def       snake_case			(					self						)		->       Optional[Any]:
														# Returns the path of the created file.
														def create_cache_and_write_probe(__UpperCAmelCase						):
																					# makedirs with exist_ok will recursively create the directory. It will not throw an error if directories
																					# already exist.
																					os.makedirs(self._cache_dir    ,						exist_ok=__UpperCAmelCase						)
																					A      :					Optional[Any]      		=					os.path.join(self._cache_dir    ,						'''fs_test''' + uuid.uuida().hex						)
																					# Opening the file in append mode will create a new file unless it already exists, in which case it will not
																					# change the file contents.
																					open(__UpperCAmelCase    ,						'''a'''						)
																					return [probe_file]

														if self._spark.conf.get('''spark.master'''    ,						''''''						).startswith('''local'''						):
																					return

														# If the cluster is multi-node, make sure that the user provided a cache_dir and that it is on an NFS
														# accessible to the driver.
														# TODO: Stream batches to the driver using ArrowCollectSerializer instead of throwing an error.
														if self._cache_dir:
																					A      :					Optional[int]      		=					(
																					    self._spark.sparkContext.parallelize(range(1						)    ,						1						).mapPartitions(__UpperCAmelCase						).collect()
																					)
																					if os.path.isfile(probe[0]						):
																												return

														raise ValueError(
														    '''When using Dataset.from_spark on a multi-node cluster, the driver and all workers should be able to access cache_dir'''						)




							def       snake_case			(					self						)		->       List[Any]:
														return datasets.DatasetInfo(features=self.config.features						)




							def       snake_case			(					self    ,						__UpperCAmelCase						)		->       str:
														return [datasets.SplitGenerator(name=datasets.Split.TRAIN						)]




							def       snake_case			(					self    ,						__UpperCAmelCase						)		->       List[Any]:
														import pyspark

														def get_arrow_batch_size(__UpperCAmelCase						):
																					for batch in it:
																												yield pa.RecordBatch.from_pydict({'''batch_bytes''': [batch.nbytes]}						)

														A      :					Any      		=					self.df.count()
														A      :					Tuple      		=					df_num_rows if df_num_rows <= 1_00 else 1_00
														# Approximate the size of each row (in Arrow format) by averaging over a max-100-row sample.
														A      :					Dict      		=					(
														    self.df.limit(__UpperCAmelCase						)
														    .repartition(1						)
														    .mapInArrow(__UpperCAmelCase    ,						'''batch_bytes: long'''						)
														    .agg(pyspark.sql.functions.sum('''batch_bytes'''						).alias('''sample_bytes'''						)						)
														    .collect()[0]
														    .sample_bytes
														    / sample_num_rows
														)
														A      :					str      		=					approx_bytes_per_row * df_num_rows
														if approx_total_size > max_shard_size:
																					# Make sure there is at least one row per partition.
																					A      :					Union[str, Any]      		=					min(__UpperCAmelCase    ,						int(approx_total_size / max_shard_size						)						)
																					A      :					List[Any]      		=					self.df.repartition(__UpperCAmelCase						)




							def       snake_case			(					self    ,						__UpperCAmelCase    ,						__UpperCAmelCase    ,						__UpperCAmelCase    ,						)		->       Iterable[Tuple[int, bool, Union[int, tuple]]]:
														import pyspark

														A      :					Dict      		=					ParquetWriter if file_format == '''parquet''' else ArrowWriter
														A      :					Optional[Any]      		=					os.path.join(self._working_dir    ,						os.path.basename(__UpperCAmelCase						)						) if self._working_dir else fpath
														A      :					int      		=					file_format == '''parquet'''

														# Define these so that we don't reference self in write_arrow, which will result in a pickling error due to
														# pickling the SparkContext.
														A      :					Any      		=					self.config.features
														A      :					Any      		=					self._writer_batch_size
														A      :					int      		=					self._fs.storage_options

														def write_arrow(__UpperCAmelCase						):
																					# Within the same SparkContext, no two task attempts will share the same attempt ID.
																					A      :					List[Any]      		=					pyspark.TaskContext().taskAttemptId()
																					A      :					Union[str, Any]      		=					next(__UpperCAmelCase    ,						__UpperCAmelCase						)
																					if first_batch is None:
																												# Some partitions might not receive any data.
																												return pa.RecordBatch.from_arrays(
																												    [[task_id], [0], [0]]    ,						names=['''task_id''', '''num_examples''', '''num_bytes''']    ,						)
																					A      :					Optional[int]      		=					0
																					A      :					Optional[int]      		=					writer_class(
																					    features=__UpperCAmelCase    ,						path=working_fpath.replace('''SSSSS'''    ,						f'{shard_id:05d}'						).replace('''TTTTT'''    ,						f'{task_id:05d}'						)    ,						writer_batch_size=__UpperCAmelCase    ,						storage_options=__UpperCAmelCase    ,						embed_local_files=__UpperCAmelCase    ,						)
																					A      :					Tuple      		=					pa.Table.from_batches([first_batch]						)
																					writer.write_table(__UpperCAmelCase						)
																					for batch in it:
																												if max_shard_size is not None and writer._num_bytes >= max_shard_size:
																																			A		,		A      :					str      		=					writer.finalize()
																																			writer.close()
																																			yield pa.RecordBatch.from_arrays(
																																			    [[task_id], [num_examples], [num_bytes]]    ,						names=['''task_id''', '''num_examples''', '''num_bytes''']    ,						)
																																			shard_id += 1
																																			A      :					List[str]      		=					writer_class(
																																			    features=writer._features    ,						path=working_fpath.replace('''SSSSS'''    ,						f'{shard_id:05d}'						).replace('''TTTTT'''    ,						f'{task_id:05d}'						)    ,						writer_batch_size=__UpperCAmelCase    ,						storage_options=__UpperCAmelCase    ,						embed_local_files=__UpperCAmelCase    ,						)
																												A      :					str      		=					pa.Table.from_batches([batch]						)
																												writer.write_table(__UpperCAmelCase						)

																					if writer._num_bytes > 0:
																												A		,		A      :					Tuple      		=					writer.finalize()
																												writer.close()
																												yield pa.RecordBatch.from_arrays(
																												    [[task_id], [num_examples], [num_bytes]]    ,						names=['''task_id''', '''num_examples''', '''num_bytes''']    ,						)

																					if working_fpath != fpath:
																												for file in os.listdir(os.path.dirname(__UpperCAmelCase						)						):
																																			A      :					Union[str, Any]      		=					os.path.join(os.path.dirname(__UpperCAmelCase						)    ,						os.path.basename(__UpperCAmelCase						)						)
																																			shutil.move(__UpperCAmelCase    ,						__UpperCAmelCase						)

														A      :					Union[str, Any]      		=					(
														    self.df.mapInArrow(__UpperCAmelCase    ,						'''task_id: long, num_examples: long, num_bytes: long'''						)
														    .groupBy('''task_id'''						)
														    .agg(
														        pyspark.sql.functions.sum('''num_examples'''						).alias('''total_num_examples'''						)    ,						pyspark.sql.functions.sum('''num_bytes'''						).alias('''total_num_bytes'''						)    ,						pyspark.sql.functions.count('''num_bytes'''						).alias('''num_shards'''						)    ,						pyspark.sql.functions.collect_list('''num_examples'''						).alias('''shard_lengths'''						)    ,						)
														    .collect()
														)
														for row in stats:
																					yield row.task_id, (row.total_num_examples, row.total_num_bytes, row.num_shards, row.shard_lengths)




							def       snake_case			(					self    ,						__UpperCAmelCase    ,						__UpperCAmelCase = "arrow"    ,						__UpperCAmelCase = None    ,						__UpperCAmelCase = None    ,						**__UpperCAmelCase    ,						)		->       List[str]:
														self._validate_cache_dir()

														A      :					int      		=					convert_file_size_to_int(max_shard_size or MAX_SHARD_SIZE						)
														self._repartition_df_if_needed(__UpperCAmelCase						)
														A      :					Any      		=					not is_remote_filesystem(self._fs						)
														A      :					Union[str, Any]      		=					os.path.join if is_local else posixpath.join

														A      :					Any      		=					'''-TTTTT-SSSSS-of-NNNNN'''
														A      :					Any      		=					f'{self.name}-{split_generator.name}{SUFFIX}.{file_format}'
														A      :					List[Any]      		=					path_join(self._output_dir    ,						__UpperCAmelCase						)

														A      :					Union[str, Any]      		=					0
														A      :					Any      		=					0
														A      :					Tuple      		=					0
														A      :					Union[str, Any]      		=					[]
														A      :					List[str]      		=					[]

														for task_id, content in self._prepare_split_single(__UpperCAmelCase    ,						__UpperCAmelCase    ,						__UpperCAmelCase						):
																					(
																					    (
																					    A
																					)		,		(
																					    A
																					)		,		(
																					    A
																					)		,		(
																					    A
																					)		,		
																					)      :					List[Any]      		=					content
																					if num_bytes > 0:
																												total_num_examples += num_examples
																												total_num_bytes += num_bytes
																												total_shards += num_shards
																												task_id_and_num_shards.append((task_id, num_shards)						)
																												all_shard_lengths.extend(__UpperCAmelCase						)

														A      :					List[str]      		=					total_num_examples
														A      :					List[str]      		=					total_num_bytes

														# should rename everything at the end
														logger.debug(f'Renaming {total_shards} shards.'						)
														if total_shards > 1:
																					A      :					Tuple      		=					all_shard_lengths

																					# Define fs outside of _rename_shard so that we don't reference self in the function, which will result in a
																					# pickling error due to pickling the SparkContext.
																					A      :					Union[str, Any]      		=					self._fs

																					# use the -SSSSS-of-NNNNN pattern
																					def _rename_shard(
																					    __UpperCAmelCase    ,						__UpperCAmelCase    ,						__UpperCAmelCase    ,						):
																												rename(
																												    __UpperCAmelCase    ,						fpath.replace('''SSSSS'''    ,						f'{shard_id:05d}'						).replace('''TTTTT'''    ,						f'{task_id:05d}'						)    ,						fpath.replace('''TTTTT-SSSSS'''    ,						f'{global_shard_id:05d}'						).replace('''NNNNN'''    ,						f'{total_shards:05d}'						)    ,						)

																					A      :					Tuple      		=					[]
																					A      :					Optional[int]      		=					0
																					for i in range(len(__UpperCAmelCase						)						):
																												A		,		A      :					Tuple      		=					task_id_and_num_shards[i]
																												for shard_id in range(__UpperCAmelCase						):
																																			args.append([task_id, shard_id, global_shard_id]						)
																																			global_shard_id += 1
																					self._spark.sparkContext.parallelize(__UpperCAmelCase    ,						len(__UpperCAmelCase						)						).map(lambda __UpperCAmelCase						: _rename_shard(*__UpperCAmelCase						)						).collect()
														else:
																					# don't use any pattern
																					A      :					int      		=					0
																					A      :					Tuple      		=					task_id_and_num_shards[0][0]
																					self._rename(
																					    fpath.replace('''SSSSS'''    ,						f'{shard_id:05d}'						).replace('''TTTTT'''    ,						f'{task_id:05d}'						)    ,						fpath.replace(__UpperCAmelCase    ,						''''''						)    ,						)






							def       snake_case			(					self    ,						__UpperCAmelCase    ,						)		->       SparkExamplesIterable:
														return SparkExamplesIterable(self.df						)






