





"""simple docstring"""




import tempfile
import unittest

import transformers
from transformers import XGLMConfig, XGLMTokenizer, is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax, require_sentencepiece, slow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


if is_flax_available():
						import jax
						import jax.numpy as jnp
						import numpy as np

						from transformers.modeling_flax_pytorch_utils import (
						    convert_pytorch_state_dict_to_flax,
						    load_flax_weights_in_pytorch_model,
						)
						from transformers.models.xglm.modeling_flax_xglm import FlaxXGLMForCausalLM, FlaxXGLMModel


if is_torch_available():
						import torch




@require_flax
class 					_a		:






				def __init__(					self						:       str,					lowerCAmelCase__						:       int,					lowerCAmelCase__						:       Optional[int]=1_4,					lowerCAmelCase__						:       List[Any]=7,					lowerCAmelCase__						:       Tuple=True,					lowerCAmelCase__						:       Any=True,					lowerCAmelCase__						:       Any=True,					lowerCAmelCase__						:       Optional[Any]=9_9,					lowerCAmelCase__						:       str=3_2,					lowerCAmelCase__						:       Union[str, Any]=5,					lowerCAmelCase__						:       Tuple=4,					lowerCAmelCase__						:       Dict=3_7,					lowerCAmelCase__						:       int="gelu",					lowerCAmelCase__						:       int=0.1,					lowerCAmelCase__						:       int=0.1,					lowerCAmelCase__						:       Union[str, Any]=5_1_2,					lowerCAmelCase__						:       Tuple=0.02,					lowerCAmelCase__						:       str=None,					)      -> Tuple:




										'''simple docstring'''



										_UpperCamelCase     :   Any      =       parent
										_UpperCamelCase     :   Union[str, Any]      =       batch_size
										_UpperCamelCase     :   List[Any]      =       seq_length
										_UpperCamelCase     :   Optional[int]      =       is_training
										_UpperCamelCase     :   List[str]      =       use_input_mask
										_UpperCamelCase     :   str      =       use_labels
										_UpperCamelCase     :   int      =       vocab_size
										_UpperCamelCase     :   str      =       d_model
										_UpperCamelCase     :   Tuple      =       num_hidden_layers
										_UpperCamelCase     :   Optional[int]      =       num_attention_heads
										_UpperCamelCase     :   int      =       ffn_dim
										_UpperCamelCase     :   Any      =       activation_function
										_UpperCamelCase     :   Tuple      =       activation_dropout
										_UpperCamelCase     :   Union[str, Any]      =       attention_dropout
										_UpperCamelCase     :   Optional[Any]      =       max_position_embeddings
										_UpperCamelCase     :   str      =       initializer_range
										_UpperCamelCase     :   Tuple      =       None
										_UpperCamelCase     :   Optional[int]      =       0
										_UpperCamelCase     :   List[str]      =       2
										_UpperCamelCase     :   Optional[Any]      =       1






				def 							snake_case      (					self						:       List[Any]				)      -> List[Any]:




										'''simple docstring'''



										_UpperCamelCase     :   str      =       np.clip(ids_tensor([self.batch_size, self.seq_length],					self.vocab_size				),					3,					self.vocab_size				)

										_UpperCamelCase     :   Any      =       None
										if self.use_input_mask:
																_UpperCamelCase     :   Union[str, Any]      =       random_attention_mask([self.batch_size, self.seq_length]				)

										_UpperCamelCase     :   int      =       XGLMConfig(
										    vocab_size=self.vocab_size,					d_model=self.hidden_size,					num_layers=self.num_hidden_layers,					attention_heads=self.num_attention_heads,					ffn_dim=self.ffn_dim,					activation_function=self.activation_function,					activation_dropout=self.activation_dropout,					attention_dropout=self.attention_dropout,					max_position_embeddings=self.max_position_embeddings,					initializer_range=self.initializer_range,					use_cache=snake_case__,					bos_token_id=self.bos_token_id,					eos_token_id=self.eos_token_id,					pad_token_id=self.pad_token_id,					)

										return (config, input_ids, input_mask)






				def 							snake_case      (					self						:       List[str]				)      -> Dict:




										'''simple docstring'''



										_UpperCamelCase     :   List[Any]      =       self.prepare_config_and_inputs()
										_UpperCamelCase  ,				_UpperCamelCase  ,				_UpperCamelCase     :   List[Any]      =       config_and_inputs
										_UpperCamelCase     :   Union[str, Any]      =       {'''input_ids''': input_ids, '''attention_mask''': attention_mask}
										return config, inputs_dict






				def 							snake_case      (					self						:       str				)      -> List[Any]:




										'''simple docstring'''



										_UpperCamelCase  ,				_UpperCamelCase  ,				_UpperCamelCase     :   Optional[int]      =       self.prepare_config_and_inputs()

										_UpperCamelCase     :   str      =       floats_tensor([self.batch_size, self.seq_length, self.hidden_size]				)
										_UpperCamelCase     :   Tuple      =       ids_tensor([self.batch_size, self.seq_length],					vocab_size=2				)

										return (
										    config,
										    input_ids,
										    attention_mask,
										    encoder_hidden_states,
										    encoder_attention_mask,
										)






				def 							snake_case      (					self						:       Tuple,					lowerCAmelCase__						:       List[str],					lowerCAmelCase__						:       str,					lowerCAmelCase__						:       Tuple,					lowerCAmelCase__						:       Optional[Any]				)      -> Optional[Any]:




										'''simple docstring'''



										_UpperCamelCase     :   str      =       2_0
										_UpperCamelCase     :   str      =       model_class_name(snake_case__				)

										_UpperCamelCase     :   Optional[int]      =       model.init_cache(input_ids.shape[0],					snake_case__				)
										_UpperCamelCase     :   Optional[int]      =       jnp.ones((input_ids.shape[0], max_decoder_length),					dtype='''i4'''				)

										_UpperCamelCase     :   Any      =       jnp.broadcast_to(
										    jnp.arange(input_ids.shape[-1] - 1				)[None, :],					(input_ids.shape[0], input_ids.shape[-1] - 1)				)
										_UpperCamelCase     :   str      =       model(
										    input_ids[:, :-1],					attention_mask=snake_case__,					past_key_values=snake_case__,					position_ids=snake_case__,					)

										_UpperCamelCase     :   Any      =       jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]],					dtype='''i4'''				)
										_UpperCamelCase     :   Tuple      =       model(
										    input_ids[:, -1:],					attention_mask=snake_case__,					past_key_values=outputs_cache.past_key_values,					position_ids=snake_case__,					)

										_UpperCamelCase     :   Optional[int]      =       model(snake_case__				)

										_UpperCamelCase     :   str      =       np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])				)				)
										self.parent.assertTrue(diff < 1e-3,					msg=f"""Max diff is {diff}"""				)






				def 							snake_case      (					self						:       Tuple,					lowerCAmelCase__						:       Union[str, Any],					lowerCAmelCase__						:       int,					lowerCAmelCase__						:       List[str],					lowerCAmelCase__						:       int				)      -> List[str]:




										'''simple docstring'''



										_UpperCamelCase     :   Optional[Any]      =       2_0
										_UpperCamelCase     :   int      =       model_class_name(snake_case__				)

										_UpperCamelCase     :   str      =       jnp.concatenate(
										    [attention_mask, jnp.zeros((attention_mask.shape[0], max_decoder_length - attention_mask.shape[1])				)],					axis=-1,					)

										_UpperCamelCase     :   str      =       model.init_cache(input_ids.shape[0],					snake_case__				)
										_UpperCamelCase     :   int      =       jnp.broadcast_to(
										    jnp.arange(input_ids.shape[-1] - 1				)[None, :],					(input_ids.shape[0], input_ids.shape[-1] - 1)				)

										_UpperCamelCase     :   Tuple      =       model(
										    input_ids[:, :-1],					attention_mask=snake_case__,					past_key_values=snake_case__,					position_ids=snake_case__,					)
										_UpperCamelCase     :   Tuple      =       jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]],					dtype='''i4'''				)
										_UpperCamelCase     :   Any      =       model(
										    input_ids[:, -1:],					past_key_values=outputs_cache.past_key_values,					attention_mask=snake_case__,					position_ids=snake_case__,					)

										_UpperCamelCase     :   Tuple      =       model(snake_case__,					attention_mask=snake_case__				)
										_UpperCamelCase     :   int      =       np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])				)				)
										self.parent.assertTrue(diff < 1e-3,					msg=f"""Max diff is {diff}"""				)




@require_sentencepiece
@require_flax
class 					_a		( a_      ,							a_      ,							unittest.TestCase							):
				UpperCamelCase								=       (FlaxXGLMModel, FlaxXGLMForCausalLM) if is_flax_available() else ()
				UpperCamelCase								=       (FlaxXGLMForCausalLM,) if is_flax_available() else ()






				def 							snake_case      (					self						:       Union[str, Any]				)      -> Any:




										'''simple docstring'''



										_UpperCamelCase     :   str      =       FlaxXGLMModelTester(self				)






				def 							snake_case      (					self						:       Dict				)      -> List[Any]:




										'''simple docstring'''



										for model_class_name in self.all_model_classes:
																_UpperCamelCase  ,				_UpperCamelCase  ,				_UpperCamelCase     :   Optional[Any]      =       self.model_tester.prepare_config_and_inputs()
																self.model_tester.check_use_cache_forward(snake_case__,					snake_case__,					snake_case__,					snake_case__				)






				def 							snake_case      (					self						:       List[Any]				)      -> Dict:




										'''simple docstring'''



										for model_class_name in self.all_model_classes:
																_UpperCamelCase  ,				_UpperCamelCase  ,				_UpperCamelCase     :   List[Any]      =       self.model_tester.prepare_config_and_inputs()
																self.model_tester.check_use_cache_forward_with_attn_mask(
																    snake_case__,					snake_case__,					snake_case__,					snake_case__				)






				@slow
				def 							snake_case      (					self						:       List[Any]				)      -> Union[str, Any]:




										'''simple docstring'''



										_UpperCamelCase     :   Union[str, Any]      =       XGLMTokenizer.from_pretrained('''XGLM''',					padding_side='''left'''				)
										_UpperCamelCase     :   Optional[int]      =       tokenizer(['''Hello this is a long string''', '''Hey'''],					return_tensors='''np''',					padding=snake_case__,					truncation=snake_case__				)

										_UpperCamelCase     :   str      =       FlaxXGLMForCausalLM.from_pretrained('''facebook/xglm-564M'''				)
										_UpperCamelCase     :   Dict      =       1
										_UpperCamelCase     :   Optional[Any]      =       False

										_UpperCamelCase     :   List[Any]      =       jax.jit(model.generate				)

										_UpperCamelCase     :   Optional[int]      =       jit_generate(inputs['''input_ids'''],					attention_mask=inputs['''attention_mask''']				).sequences

										_UpperCamelCase     :   int      =       tokenizer.batch_decode(snake_case__,					skip_special_tokens=snake_case__				)

										_UpperCamelCase     :   str      =       [
										    '''Hello this is a long string of questions, but I\'m not sure if I\'m''',
										    '''Hey, I\'m a newbie to the forum and I\'''',
										]

										self.assertListEqual(snake_case__,					snake_case__				)






				@is_pt_flax_cross_test
				def 							snake_case      (					self						:       int				)      -> Union[str, Any]:




										'''simple docstring'''



										_UpperCamelCase  ,				_UpperCamelCase     :   Any      =       self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes:
																with self.subTest(model_class.__name__				):
																						# prepare inputs
																						_UpperCamelCase     :   List[str]      =       self._prepare_for_class(snake_case__,					snake_case__				)
																						_UpperCamelCase     :   Optional[Any]      =       {k: torch.tensor(v.tolist()				) for k, v in prepared_inputs_dict.items()}

																						# load corresponding PyTorch class
																						_UpperCamelCase     :   Dict      =       model_class.__name__[4:]  # Skip the "Flax" at the beginning
																						_UpperCamelCase     :   Dict      =       getattr(snake_case__,					snake_case__				)

																						_UpperCamelCase  ,				_UpperCamelCase     :   Tuple      =       pt_inputs['''input_ids'''].shape
																						_UpperCamelCase     :   List[str]      =       np.random.randint(0,					seq_length - 1,					size=(batch_size,)				)
																						for batch_idx, start_index in enumerate(snake_case__				):
																												_UpperCamelCase     :   List[Any]      =       0
																												_UpperCamelCase     :   List[str]      =       1
																												_UpperCamelCase     :   List[Any]      =       0
																												_UpperCamelCase     :   Any      =       1
																						_UpperCamelCase     :   Tuple      =       pt_model_class(snake_case__				).eval()
																						# Flax models don't use the `use_cache` option and cache is not returned as a default.
																						# So we disable `use_cache` here for PyTorch model.
																						_UpperCamelCase     :   Dict      =       False
																						_UpperCamelCase     :   List[str]      =       model_class(snake_case__,					dtype=jnp.floataa				)
																						_UpperCamelCase     :   Dict      =       convert_pytorch_state_dict_to_flax(pt_model.state_dict(),					snake_case__				)
																						_UpperCamelCase     :   Any      =       fx_state

																						with torch.no_grad():
																												_UpperCamelCase     :   Union[str, Any]      =       pt_model(**snake_case__				).to_tuple()

																						_UpperCamelCase     :   Optional[Any]      =       fx_model(**snake_case__				).to_tuple()
																						self.assertEqual(len(snake_case__				),					len(snake_case__				),					'''Output lengths differ between Flax and PyTorch'''				)
																						for fx_output, pt_output in zip(snake_case__,					snake_case__				):
																												self.assert_almost_equals(fx_output[:, -1],					pt_output[:, -1].numpy(),					4e-2				)

																						with tempfile.TemporaryDirectory() as tmpdirname:
																												pt_model.save_pretrained(snake_case__				)
																												_UpperCamelCase     :   Optional[int]      =       model_class.from_pretrained(snake_case__,					from_pt=snake_case__				)

																						_UpperCamelCase     :   Optional[int]      =       fx_model_loaded(**snake_case__				).to_tuple()
																						self.assertEqual(
																						    len(snake_case__				),					len(snake_case__				),					'''Output lengths differ between Flax and PyTorch'''				)
																						for fx_output_loaded, pt_output in zip(snake_case__,					snake_case__				):
																												self.assert_almost_equals(fx_output_loaded[:, -1],					pt_output[:, -1].numpy(),					4e-2				)






				@is_pt_flax_cross_test
				def 							snake_case      (					self						:       Union[str, Any]				)      -> Optional[int]:




										'''simple docstring'''



										_UpperCamelCase  ,				_UpperCamelCase     :   Optional[Any]      =       self.model_tester.prepare_config_and_inputs_for_common()
										for model_class in self.all_model_classes:
																with self.subTest(model_class.__name__				):
																						# prepare inputs
																						_UpperCamelCase     :   Tuple      =       self._prepare_for_class(snake_case__,					snake_case__				)
																						_UpperCamelCase     :   Union[str, Any]      =       {k: torch.tensor(v.tolist()				) for k, v in prepared_inputs_dict.items()}

																						# load corresponding PyTorch class
																						_UpperCamelCase     :   List[Any]      =       model_class.__name__[4:]  # Skip the "Flax" at the beginning
																						_UpperCamelCase     :   List[str]      =       getattr(snake_case__,					snake_case__				)

																						_UpperCamelCase     :   Tuple      =       pt_model_class(snake_case__				).eval()
																						_UpperCamelCase     :   Optional[int]      =       False
																						_UpperCamelCase     :   int      =       model_class(snake_case__,					dtype=jnp.floataa				)

																						_UpperCamelCase     :   Optional[int]      =       load_flax_weights_in_pytorch_model(snake_case__,					fx_model.params				)
																						_UpperCamelCase  ,				_UpperCamelCase     :   Tuple      =       pt_inputs['''input_ids'''].shape
																						_UpperCamelCase     :   Union[str, Any]      =       np.random.randint(0,					seq_length - 1,					size=(batch_size,)				)
																						for batch_idx, start_index in enumerate(snake_case__				):
																												_UpperCamelCase     :   str      =       0
																												_UpperCamelCase     :   List[str]      =       1
																												_UpperCamelCase     :   Dict      =       0
																												_UpperCamelCase     :   Optional[Any]      =       1

																						# make sure weights are tied in PyTorch
																						pt_model.tie_weights()

																						with torch.no_grad():
																												_UpperCamelCase     :   Optional[Any]      =       pt_model(**snake_case__				).to_tuple()

																						_UpperCamelCase     :   List[Any]      =       fx_model(**snake_case__				).to_tuple()
																						self.assertEqual(len(snake_case__				),					len(snake_case__				),					'''Output lengths differ between Flax and PyTorch'''				)
																						for fx_output, pt_output in zip(snake_case__,					snake_case__				):
																												self.assert_almost_equals(fx_output[:, -1],					pt_output[:, -1].numpy(),					4e-2				)

																						with tempfile.TemporaryDirectory() as tmpdirname:
																												fx_model.save_pretrained(snake_case__				)
																												_UpperCamelCase     :   Union[str, Any]      =       pt_model_class.from_pretrained(snake_case__,					from_flax=snake_case__				)

																						with torch.no_grad():
																												_UpperCamelCase     :   List[Any]      =       pt_model_loaded(**snake_case__				).to_tuple()

																						self.assertEqual(
																						    len(snake_case__				),					len(snake_case__				),					'''Output lengths differ between Flax and PyTorch'''				)
																						for fx_output, pt_output in zip(snake_case__,					snake_case__				):
																												self.assert_almost_equals(fx_output[:, -1],					pt_output[:, -1].numpy(),					4e-2				)






				@slow
				def 							snake_case      (					self						:       List[str]				)      -> Dict:




										'''simple docstring'''



										for model_class_name in self.all_model_classes:
																_UpperCamelCase     :   List[str]      =       model_class_name.from_pretrained('''facebook/xglm-564M'''				)
																_UpperCamelCase     :   Any      =       model(np.ones((1, 1)				)				)
																self.assertIsNotNone(snake_case__				)








import json
import os
from dataclasses import dataclass

import numpy as np
import pyarrow as pa

import datasets
from utils import get_duration


UpperCAmelCase :				Any							=1000_0000_0000
UpperCAmelCase :				Any							=1_0000

UpperCAmelCase			,			UpperCAmelCase :				Tuple							=os.path.split(__file__)
UpperCAmelCase :				Tuple							=os.path.join(RESULTS_BASEPATH, """results""", RESULTS_FILENAME.replace(""".py""", """.json"""))



def  _lowerCAmelCase     (_lowerCAmelCase ,						_lowerCAmelCase):
					UpperCamelCase_			= pa.Table.from_pydict({"col": [0] * chunk_size})
					UpperCamelCase_			= pa.concat_tables([table] * (num_examples // chunk_size))
					return datasets.Dataset(_lowerCAmelCase ,						fingerprint="table_100B")


@dataclass
class    _lowercase :
				'''simple docstring'''







				lowercase__     =    42
				lowercase__     =    42
				lowercase__     =    42
				lowercase__     =    42

				def    _lowerCamelCase       (			self     ):

									'''simple docstring'''


									UpperCamelCase_			= np.random.default_rng(self.seed     )
									UpperCamelCase_			= rng.integers(low=self.low  ,			high=self.high  ,			size=self.size     ).tolist()

				def __iter__(			self     ):

									'''simple docstring'''


									return iter(self._sampled_values     )

				def __len__(			self     ):

									'''simple docstring'''


									return self.size




@get_duration
def  _lowerCAmelCase     (_lowerCAmelCase):
					UpperCamelCase_			= dataset[0]



@get_duration
def  _lowerCAmelCase     (_lowerCAmelCase):
					UpperCamelCase_			= dataset[-1]



@get_duration
def  _lowerCAmelCase     (_lowerCAmelCase):
					UpperCamelCase_			= dataset[range(len(_lowerCAmelCase) // 2 ,						len(_lowerCAmelCase) // 2 + 10_24)]



@get_duration
def  _lowerCAmelCase     (_lowerCAmelCase):
					UpperCamelCase_			= dataset[RandIter(0 ,						len(_lowerCAmelCase) ,						10_24 ,						seed=42)]



def  _lowerCAmelCase     ():
					UpperCamelCase_			= {"num examples": SPEED_TEST_N_EXAMPLES}
					UpperCamelCase_			= (get_first_row, get_last_row, get_batch_of_aaaa_rows, get_batch_of_aaaa_random_rows)
					print("generating dataset")
					UpperCamelCase_			= generate_100B_dataset(num_examples=_lowerCAmelCase ,						chunk_size=_lowerCAmelCase)
					print("Functions")
					for func in functions:
										print(func.__name__)
										UpperCamelCase_			= func(_lowerCAmelCase)

					with open(_lowerCAmelCase ,						"wb") as f:
										f.write(json.dumps(_lowerCAmelCase).encode("utf-8"))


if __name__ == "__main__":  # useful to run the profiler
	benchmark_table_100B()
