
import unittest

from transformers import BigBirdConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
						import jax

						from transformers.models.big_bird.modeling_flax_big_bird import (
						    FlaxBigBirdForCausalLM,
						    FlaxBigBirdForMaskedLM,
						    FlaxBigBirdForMultipleChoice,
						    FlaxBigBirdForPreTraining,
						    FlaxBigBirdForQuestionAnswering,
						    FlaxBigBirdForSequenceClassification,
						    FlaxBigBirdForTokenClassification,
						    FlaxBigBirdModel,
						)
class  lowercase__(				unittest.TestCase							):



	"""simple docstring"""




	def __init__(       self  :			Optional[int]       ,							SCREAMING_SNAKE_CASE_  :			List[Any]       ,							SCREAMING_SNAKE_CASE_  :			str=2       ,							SCREAMING_SNAKE_CASE_  :			Dict=5_6       ,							SCREAMING_SNAKE_CASE_  :			Union[str, Any]=True       ,							SCREAMING_SNAKE_CASE_  :			Union[str, Any]=True       ,							SCREAMING_SNAKE_CASE_  :			Any=True       ,							SCREAMING_SNAKE_CASE_  :			str=True       ,							SCREAMING_SNAKE_CASE_  :			Any=9_9       ,							SCREAMING_SNAKE_CASE_  :			Union[str, Any]=3_2       ,							SCREAMING_SNAKE_CASE_  :			Union[str, Any]=2       ,							SCREAMING_SNAKE_CASE_  :			int=2       ,							SCREAMING_SNAKE_CASE_  :			List[Any]=7       ,							SCREAMING_SNAKE_CASE_  :			Optional[Any]="gelu_new"       ,							SCREAMING_SNAKE_CASE_  :			str=0.1       ,							SCREAMING_SNAKE_CASE_  :			int=0.1       ,							SCREAMING_SNAKE_CASE_  :			Optional[int]=5_1_2       ,							SCREAMING_SNAKE_CASE_  :			int=1_6       ,							SCREAMING_SNAKE_CASE_  :			Optional[Any]=2       ,							SCREAMING_SNAKE_CASE_  :			Any=0.02       ,							SCREAMING_SNAKE_CASE_  :			Optional[Any]=4       ,							SCREAMING_SNAKE_CASE_  :			str="block_sparse"       ,							SCREAMING_SNAKE_CASE_  :			Optional[Any]=True       ,							SCREAMING_SNAKE_CASE_  :			Optional[int]=False       ,							SCREAMING_SNAKE_CASE_  :			List[Any]=2       ,							SCREAMING_SNAKE_CASE_  :			Dict=3       ,							)       ->     Union[str, Any]:
					lowercase_       =					parent
					lowercase_       =					batch_size
					lowercase_       =					seq_length
					lowercase_       =					is_training
					lowercase_       =					use_attention_mask
					lowercase_       =					use_token_type_ids
					lowercase_       =					use_labels
					lowercase_       =					vocab_size
					lowercase_       =					hidden_size
					lowercase_       =					num_hidden_layers
					lowercase_       =					num_attention_heads
					lowercase_       =					intermediate_size
					lowercase_       =					hidden_act
					lowercase_       =					hidden_dropout_prob
					lowercase_       =					attention_probs_dropout_prob
					lowercase_       =					max_position_embeddings
					lowercase_       =					type_vocab_size
					lowercase_       =					type_sequence_label_size
					lowercase_       =					initializer_range
					lowercase_       =					num_choices

					lowercase_       =					rescale_embeddings
					lowercase_       =					attention_type
					lowercase_       =					use_bias
					lowercase_       =					block_size
					lowercase_       =					num_random_blocks
	def 						_lowercase  (       self  :			Dict       )       ->     Optional[int]:
					lowercase_       =					ids_tensor([self.batch_size, self.seq_length]       ,							self.vocab_size       )

					lowercase_       =					None
					if self.use_attention_mask:
									lowercase_       =					random_attention_mask([self.batch_size, self.seq_length]       )

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

					lowercase_       =					BigBirdConfig(
					    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=SCREAMING_SNAKE_CASE_       ,							initializer_range=self.initializer_range       ,							attention_type=self.attention_type       ,							block_size=self.block_size       ,							num_random_blocks=self.num_random_blocks       ,							use_bias=self.use_bias       ,							rescale_embeddings=self.rescale_embeddings       ,							)

					return config, input_ids, token_type_ids, attention_mask




	def 						_lowercase  (       self  :			Tuple       )       ->     Dict:
					lowercase_       =					self.prepare_config_and_inputs()
					lowercase_       =					config_and_inputs
					lowercase_       =					{
					    '''input_ids''': input_ids,
					    '''token_type_ids''': token_type_ids,
					    '''attention_mask''': attention_mask,
					}
					return config, inputs_dict

@require_flax
class  lowercase__(				_a ,						unittest.TestCase							):



	"""simple docstring"""




	a			:str         =	(
	    (
	        FlaxBigBirdForCausalLM,
	        FlaxBigBirdModel,
	        FlaxBigBirdForPreTraining,
	        FlaxBigBirdForMaskedLM,
	        FlaxBigBirdForMultipleChoice,
	        FlaxBigBirdForQuestionAnswering,
	        FlaxBigBirdForSequenceClassification,
	        FlaxBigBirdForTokenClassification,
	    )
	    if is_flax_available()
	    else ()
	)

	a			:Any         =	False
	a			:List[str]         =	False
	def 						_lowercase  (       self  :			int       )       ->     Union[str, Any]:
					lowercase_       =					FlaxBigBirdModelTester(self       )
	@slow
	# copied from `test_modeling_flax_common` because it takes much longer than other models
	def 						_lowercase  (       self  :			Any       )       ->     Dict:
					super().test_from_pretrained_save_pretrained()
	@slow
	# copied from `test_modeling_flax_common` because it takes much longer than other models
	def 						_lowercase  (       self  :			List[str]       )       ->     List[str]:
					super().test_from_pretrained_with_no_automatic_init()
	@slow
	# copied from `test_modeling_flax_common` because it takes much longer than other models
	def 						_lowercase  (       self  :			List[Any]       )       ->     Tuple:
					super().test_no_automatic_init()
	@slow
	# copied from `test_modeling_flax_common` because it takes much longer than other models
	def 						_lowercase  (       self  :			Tuple       )       ->     Any:
					super().test_hidden_states_output()
	@slow
	def 						_lowercase  (       self  :			List[Any]       )       ->     Dict:



					for model_class_name in self.all_model_classes:
									lowercase_       =					model_class_name.from_pretrained('''google/bigbird-roberta-base'''       )
									self.assertIsNotNone(SCREAMING_SNAKE_CASE_       )
	def 						_lowercase  (       self  :			List[Any]       )       ->     int:



					if self.test_attn_probs:
									super().test_attention_outputs()
	@slow
	# copied from `test_modeling_flax_common` because it takes much longer than other models
	def 						_lowercase  (       self  :			Tuple       )       ->     Tuple:
					lowercase_       =					self.model_tester.prepare_config_and_inputs_for_common()

					for model_class in self.all_model_classes:
									with self.subTest(model_class.__name__       ):
													lowercase_       =					self._prepare_for_class(SCREAMING_SNAKE_CASE_       ,							SCREAMING_SNAKE_CASE_       )
													lowercase_       =					model_class(SCREAMING_SNAKE_CASE_       )

													@jax.jit
													def model_jitted(SCREAMING_SNAKE_CASE_  :			Dict       ,							SCREAMING_SNAKE_CASE_  :			Optional[Any]=None       ,							**SCREAMING_SNAKE_CASE_  :			List[Any]       ):
																	return model(input_ids=SCREAMING_SNAKE_CASE_       ,							attention_mask=SCREAMING_SNAKE_CASE_       ,							**SCREAMING_SNAKE_CASE_       )

													with self.subTest('''JIT Enabled'''       ):
																	lowercase_       =					model_jitted(**SCREAMING_SNAKE_CASE_       ).to_tuple()

													with self.subTest('''JIT Disabled'''       ):
																	with jax.disable_jit():
																					lowercase_       =					model_jitted(**SCREAMING_SNAKE_CASE_       ).to_tuple()

													self.assertEqual(len(SCREAMING_SNAKE_CASE_       )       ,							len(SCREAMING_SNAKE_CASE_       )       )
													for jitted_output, output in zip(SCREAMING_SNAKE_CASE_       ,							SCREAMING_SNAKE_CASE_       ):
																	self.assertEqual(jitted_output.shape       ,							output.shape       )




	def 						_lowercase  (       self  :			int       ,							SCREAMING_SNAKE_CASE_  :			str       ,							SCREAMING_SNAKE_CASE_  :			Dict       ,							SCREAMING_SNAKE_CASE_  :			Optional[int]       ,							SCREAMING_SNAKE_CASE_  :			Dict=1e-5       ,							SCREAMING_SNAKE_CASE_  :			Dict="outputs"       ,							SCREAMING_SNAKE_CASE_  :			Optional[int]=None       )       ->     List[str]:



					if name.startswith('''outputs.attentions'''       ):
									return
					else:
									super().check_pt_flax_outputs(SCREAMING_SNAKE_CASE_       ,							SCREAMING_SNAKE_CASE_       ,							SCREAMING_SNAKE_CASE_       ,							SCREAMING_SNAKE_CASE_       ,							SCREAMING_SNAKE_CASE_       ,							SCREAMING_SNAKE_CASE_       )











import argparse
import logging
import os

import datasets
import tensorflow as tf

from transformers import AutoTokenizer


lowerCAmelCase_    =				logging.getLogger(__name__)






def      lowerCamelCase_						(				)		->						Optional[Any]:






				"""simple docstring"""


				snake_case_						:    List[str]				=  argparse.ArgumentParser(
				    description='''Prepare TFRecord shards from pre-tokenized samples of the wikitext dataset.''' )
				parser.add_argument(
				    '''--dataset_name''' ,				type=_UpperCamelCase ,				default='''wikitext''' ,				help='''Name of the training. Explore datasets at: hf.co/datasets.''' ,				)
				parser.add_argument(
				    '''--dataset_config''' ,				type=_UpperCamelCase ,				default='''wikitext-103-raw-v1''' ,				help='''Configuration name of the dataset.''' )
				parser.add_argument(
				    '''--tokenizer_name_or_path''' ,				type=_UpperCamelCase ,				default='''sayakpaul/unigram-tokenizer-wikitext''' ,				help='''Tokenizer identifier. Can be a local filepath or a Hub identifier.''' ,				)
				parser.add_argument(
				    '''--shard_size''' ,				type=_UpperCamelCase ,				default=1_000 ,				help='''Number of entries to go in a single shard.''' ,				)
				parser.add_argument('''--split''' ,				type=_UpperCamelCase ,				default='''train''' ,				choices=['''train''', '''test''', '''validation'''] )
				parser.add_argument(
				    '''--limit''' ,				default=_UpperCamelCase ,				type=_UpperCamelCase ,				help='''Limit the number of shards (used for debugging).''' ,				)
				parser.add_argument(
				    '''--max_length''' ,				type=_UpperCamelCase ,				default=512 ,				help='''Maximum sequence length. For training on TPUs, it helps to have a maximum'''
				    ''' sequence length that is a multiple of 8.''' ,				)
				parser.add_argument(
				    '''--output_dir''' ,				default='''tf-tpu''' ,				type=_UpperCamelCase ,				help='''Output directory where the TFRecord shards will be saved. If the'''
				    ''' path is appended with `gs://` (\'gs://tf-tpu\', for example) then the TFRecord'''
				    ''' shards will be directly saved to a Google Cloud Storage bucket.''' ,				)

				snake_case_						:    List[Any]				=  parser.parse_args()
				return args






def      lowerCamelCase_						(				_UpperCamelCase )		->						Tuple:






				"""simple docstring"""


				def fn(_UpperCamelCase ):
								return tokenizer(examples['''text'''] )

				return fn






def      lowerCamelCase_						(				_UpperCamelCase )		->						Union[str, Any]:






				"""simple docstring"""


				snake_case_						:    Any				=  []
				for i in range(len(tokenized_data['''input_ids'''] ) ):
								snake_case_						:    Any				=  {
								    '''input_ids''': tf.train.Feature(intaa_list=tf.train.IntaaList(value=tokenized_data['''input_ids'''][i] ) ),
								    '''attention_mask''': tf.train.Feature(
								        intaa_list=tf.train.IntaaList(value=tokenized_data['''attention_mask'''][i] ) ),
								}
								snake_case_						:    Optional[int]				=  tf.train.Features(feature=_UpperCamelCase )
								snake_case_						:    Optional[Any]				=  tf.train.Example(features=_UpperCamelCase )
								snake_case_						:    Optional[Any]				=  example.SerializeToString()
								records.append(_UpperCamelCase )
				return records






def      lowerCamelCase_						(				_UpperCamelCase )		->						Optional[int]:






				"""simple docstring"""


				snake_case_						:    int				=  datasets.load_dataset(args.dataset_name ,				args.dataset_config ,				split=args.split )

				if args.limit is not None:
								snake_case_						:    Union[str, Any]				=  min(len(_UpperCamelCase ) ,				args.limit )
								snake_case_						:    int				=  dataset.select(range(_UpperCamelCase ) )
								print(f'''Limiting the dataset to {args.limit} entries.''' )

				snake_case_						:    Dict				=  AutoTokenizer.from_pretrained(args.tokenizer_name_or_path )

				# Handle output directory creation.
				# For serializing into a Google Cloud Storage Bucket, one needs to first
				# create a bucket.
				if "gs" not in args.output_dir:
								if not os.path.exists(args.output_dir ):
												os.makedirs(args.output_dir )
								snake_case_						:    str				=  os.path.join(args.output_dir ,				args.split )
								if not os.path.exists(_UpperCamelCase ):
												os.makedirs(_UpperCamelCase )
				else:
								snake_case_						:    Optional[Any]				=  os.path.join(args.output_dir ,				args.split )

				# Tokenize the whole dataset at once.
				snake_case_						:    Optional[Any]				=  tokenize_function(_UpperCamelCase )
				snake_case_						:    List[Any]				=  dataset.map(_UpperCamelCase ,				batched=_UpperCamelCase ,				num_proc=4 ,				remove_columns=['''text'''] )

				# We need to concatenate all our texts together, and then split the result
				# into chunks of a fixed size, which we will call block_size. To do this, we
				# will use the map method again, with the option batched=True. When we use batched=True,
				# the function we pass to map() will be passed multiple inputs at once, allowing us
				# to group them into more or fewer examples than we had in the input.
				# This allows us to create our new fixed-length samples. The advantage of this
				# method is that we don't lose a whole lot of content from the dataset compared to the
				# case where we simply tokenize with a pre-defined max_length.

				def group_texts(_UpperCamelCase ):
								# Concatenate all texts.
								snake_case_						:    Tuple				=  {k: sum(examples[k] ,				[] ) for k in examples.keys()}
								snake_case_						:    List[str]				=  len(concatenated_examples[list(examples.keys() )[0]] )
								# We drop the small remainder, though you could add padding instead if the model supports it
								# In this, as in all things, we advise you to follow your heart 🫀
								snake_case_						:    int				=  (total_length // args.max_length) * args.max_length
								# Split by chunks of max_len.
								snake_case_						:    Union[str, Any]				=  {
								    k: [t[i : i + args.max_length] for i in range(0 ,				_UpperCamelCase ,				args.max_length )]
								    for k, t in concatenated_examples.items()
								}
								return result

				snake_case_						:    int				=  dataset_tokenized.map(_UpperCamelCase ,				batched=_UpperCamelCase ,				batch_size=1_000 ,				num_proc=4 )

				snake_case_						:    str				=  0
				snake_case_						:    Optional[Any]				=  0
				for shard in range(0 ,				len(_UpperCamelCase ) ,				args.shard_size ):
								snake_case_						:    Any				=  grouped_dataset[shard : shard + args.shard_size]
								snake_case_						:    str				=  len(dataset_snapshot['''input_ids'''] )
								snake_case_						:    Union[str, Any]				=  os.path.join(_UpperCamelCase ,				f'''dataset-{shard_count}-{records_containing}.tfrecord''' )
								snake_case_						:    Dict				=  get_serialized_examples(_UpperCamelCase )

								with tf.io.TFRecordWriter(_UpperCamelCase ) as out_file:
												for i in range(len(_UpperCamelCase ) ):
																snake_case_						:    List[str]				=  serialized_examples[i]
																out_file.write(_UpperCamelCase )
												print('''Wrote file {} containing {} records'''.format(_UpperCamelCase ,				_UpperCamelCase ) )

								shard_count += 1
								total_records += records_containing

				with open(f'''split-{args.split}-records-count.txt''' ,				'''w''' ) as f:
								print(f'''Total {args.split} records: {total_records}''' ,				file=_UpperCamelCase )


if __name__ == "__main__":
	lowerCAmelCase_    =				parse_args()
	main(args)
