


import argparse
import logging
import os

import datasets
import tensorflow as tf

from transformers import AutoTokenizer


__snake_case            =   logging.getLogger(__name__)






def 							_lowercase    (		):





			"""simple docstring"""




			UpperCamelCase         =					argparse.ArgumentParser(
			    description="""Prepare TFRecord shards from pre-tokenized samples of the wikitext dataset."""	)
			parser.add_argument(
			    """--dataset_name"""       ,  type=SCREAMING_SNAKE_CASE_       ,  default="""wikitext"""       ,  help="""Name of the training. Explore datasets at: hf.co/datasets."""       ,  )
			parser.add_argument(
			    """--dataset_config"""       ,  type=SCREAMING_SNAKE_CASE_       ,  default="""wikitext-103-raw-v1"""       ,  help="""Configuration name of the dataset."""	)
			parser.add_argument(
			    """--tokenizer_name_or_path"""       ,  type=SCREAMING_SNAKE_CASE_       ,  default="""sayakpaul/unigram-tokenizer-wikitext"""       ,  help="""Tokenizer identifier. Can be a local filepath or a Hub identifier."""       ,  )
			parser.add_argument(
			    """--shard_size"""       ,  type=SCREAMING_SNAKE_CASE_       ,  default=1_000       ,  help="""Number of entries to go in a single shard."""       ,  )
			parser.add_argument("""--split"""       ,  type=SCREAMING_SNAKE_CASE_       ,  default="""train"""       ,  choices=["""train""", """test""", """validation"""]	)
			parser.add_argument(
			    """--limit"""       ,  default=SCREAMING_SNAKE_CASE_       ,  type=SCREAMING_SNAKE_CASE_       ,  help="""Limit the number of shards (used for debugging)."""       ,  )
			parser.add_argument(
			    """--max_length"""       ,  type=SCREAMING_SNAKE_CASE_       ,  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=SCREAMING_SNAKE_CASE_       ,  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."""       ,  )

			UpperCamelCase         =					parser.parse_args()
			return args






def 							_lowercase    (		SCREAMING_SNAKE_CASE_   :		str	):





			"""simple docstring"""




			def fn(SCREAMING_SNAKE_CASE_   :		int	):
						return tokenizer(examples["""text"""]	)

			return fn






def 							_lowercase    (		SCREAMING_SNAKE_CASE_   :		int	):





			"""simple docstring"""




			UpperCamelCase         =					[]
			for i in range(len(tokenized_data["""input_ids"""]	)	):
						UpperCamelCase         =					{
						    """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]	)	),
						}
						UpperCamelCase         =					tf.train.Features(feature=SCREAMING_SNAKE_CASE_	)
						UpperCamelCase         =					tf.train.Example(features=SCREAMING_SNAKE_CASE_	)
						UpperCamelCase         =					example.SerializeToString()
						records.append(SCREAMING_SNAKE_CASE_	)
			return records






def 							_lowercase    (		SCREAMING_SNAKE_CASE_   :		List[Any]	):





			"""simple docstring"""




			UpperCamelCase         =					datasets.load_dataset(args.dataset_name       ,  args.dataset_config       ,  split=args.split	)

			if args.limit is not None:
						UpperCamelCase         =					min(len(SCREAMING_SNAKE_CASE_	)       ,  args.limit	)
						UpperCamelCase         =					dataset.select(range(SCREAMING_SNAKE_CASE_	)	)
						print(f'Limiting the dataset to {args.limit} entries.'	)

			UpperCamelCase         =					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	)
						UpperCamelCase         =					os.path.join(args.output_dir       ,  args.split	)
						if not os.path.exists(SCREAMING_SNAKE_CASE_	):
									os.makedirs(SCREAMING_SNAKE_CASE_	)
			else:
						UpperCamelCase         =					os.path.join(args.output_dir       ,  args.split	)

			# Tokenize the whole dataset at once.
			UpperCamelCase         =					tokenize_function(SCREAMING_SNAKE_CASE_	)
			UpperCamelCase         =					dataset.map(SCREAMING_SNAKE_CASE_       ,  batched=SCREAMING_SNAKE_CASE_       ,  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(SCREAMING_SNAKE_CASE_   :		Union[str, Any]	):
						# Concatenate all texts.
						UpperCamelCase         =					{k: sum(examples[k]       ,  []	) for k in examples.keys()}
						UpperCamelCase         =					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 🫀
						UpperCamelCase         =					(total_length // args.max_length) * args.max_length
						# Split by chunks of max_len.
						UpperCamelCase         =					{
						    k: [t[i : i + args.max_length] for i in range(0       ,  SCREAMING_SNAKE_CASE_       ,  args.max_length	)]
						    for k, t in concatenated_examples.items()
						}
						return result

			UpperCamelCase         =					dataset_tokenized.map(SCREAMING_SNAKE_CASE_       ,  batched=SCREAMING_SNAKE_CASE_       ,  batch_size=1_000       ,  num_proc=4	)

			UpperCamelCase         =					0
			UpperCamelCase         =					0
			for shard in range(0       ,  len(SCREAMING_SNAKE_CASE_	)       ,  args.shard_size	):
						UpperCamelCase         =					grouped_dataset[shard : shard + args.shard_size]
						UpperCamelCase         =					len(dataset_snapshot["""input_ids"""]	)
						UpperCamelCase         =					os.path.join(SCREAMING_SNAKE_CASE_       ,  f'dataset-{shard_count}-{records_containing}.tfrecord'	)
						UpperCamelCase         =					get_serialized_examples(SCREAMING_SNAKE_CASE_	)

						with tf.io.TFRecordWriter(SCREAMING_SNAKE_CASE_	) as out_file:
									for i in range(len(SCREAMING_SNAKE_CASE_	)	):
												UpperCamelCase         =					serialized_examples[i]
												out_file.write(SCREAMING_SNAKE_CASE_	)
									print("""Wrote file {} containing {} records""".format(SCREAMING_SNAKE_CASE_       ,  SCREAMING_SNAKE_CASE_	)	)

						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=SCREAMING_SNAKE_CASE_	)


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








import unittest

from transformers import BertGenerationConfig, is_torch_available
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, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
			import torch

			from transformers import BertGenerationDecoder, BertGenerationEncoder




class    UpperCAmelCase					:




			def __init__(						self							:		Union[str, Any]       ,       __magic_name__							:		List[Any]       ,       __magic_name__							:		str=1_3       ,       __magic_name__							:		str=7       ,       __magic_name__							:		Optional[Any]=True       ,       __magic_name__							:		Optional[Any]=True       ,       __magic_name__							:		Tuple=9_9       ,       __magic_name__							:		List[str]=3_2       ,       __magic_name__							:		List[str]=5       ,       __magic_name__							:		int=4       ,       __magic_name__							:		Union[str, Any]=3_7       ,       __magic_name__							:		List[Any]="gelu"       ,       __magic_name__							:		Any=0.1       ,       __magic_name__							:		Union[str, Any]=0.1       ,       __magic_name__							:		List[Any]=5_0       ,       __magic_name__							:		Tuple=0.02       ,       __magic_name__							:		Optional[int]=True       ,       __magic_name__							:		Optional[Any]=None       ,       ):







						"""simple docstring"""







						UpperCamelCase         =					parent
						UpperCamelCase         =					batch_size
						UpperCamelCase         =					seq_length
						UpperCamelCase         =					is_training
						UpperCamelCase         =					use_input_mask
						UpperCamelCase         =					vocab_size
						UpperCamelCase         =					hidden_size
						UpperCamelCase         =					num_hidden_layers
						UpperCamelCase         =					num_attention_heads
						UpperCamelCase         =					intermediate_size
						UpperCamelCase         =					hidden_act
						UpperCamelCase         =					hidden_dropout_prob
						UpperCamelCase         =					attention_probs_dropout_prob
						UpperCamelCase         =					max_position_embeddings
						UpperCamelCase         =					initializer_range
						UpperCamelCase         =					use_labels
						UpperCamelCase         =					scope




			def   lowerCamelCase_					(						self							:		Optional[int]						):







						"""simple docstring"""







						UpperCamelCase         =					ids_tensor([self.batch_size, self.seq_length]       ,       self.vocab_size						)

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

						if self.use_labels:
									UpperCamelCase         =					ids_tensor([self.batch_size, self.seq_length]       ,       self.vocab_size						)

						UpperCamelCase         =					self.get_config()

						return config, input_ids, input_mask, token_labels




			def   lowerCamelCase_					(						self							:		str						):







						"""simple docstring"""







						return BertGenerationConfig(
						    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       ,       is_decoder=__magic_name__       ,       initializer_range=self.initializer_range       ,       )




			def   lowerCamelCase_					(						self							:		Union[str, Any]						):







						"""simple docstring"""







						(
						    (
						    UpperCamelCase
						)  ,  (
						    UpperCamelCase
						)  ,  (
						    UpperCamelCase
						)  ,  (
						    UpperCamelCase
						)  ,  
						)         =					self.prepare_config_and_inputs()

						UpperCamelCase         =					True
						UpperCamelCase         =					floats_tensor([self.batch_size, self.seq_length, self.hidden_size]						)
						UpperCamelCase         =					ids_tensor([self.batch_size, self.seq_length]       ,       vocab_size=2						)

						return (
						    config,
						    input_ids,
						    input_mask,
						    token_labels,
						    encoder_hidden_states,
						    encoder_attention_mask,
						)




			def   lowerCamelCase_					(						self							:		Tuple       ,       __magic_name__							:		List[Any]       ,       __magic_name__							:		Tuple       ,       __magic_name__							:		str       ,       __magic_name__							:		List[str]       ,       **__magic_name__							:		Optional[Any]       ,       ):







						"""simple docstring"""







						UpperCamelCase         =					BertGenerationEncoder(config=__magic_name__						)
						model.to(__magic_name__						)
						model.eval()
						UpperCamelCase         =					model(__magic_name__       ,       attention_mask=__magic_name__						)
						UpperCamelCase         =					model(__magic_name__						)
						self.parent.assertEqual(result.last_hidden_state.shape       ,       (self.batch_size, self.seq_length, self.hidden_size)						)




			def   lowerCamelCase_					(						self							:		Optional[int]       ,       __magic_name__							:		Dict       ,       __magic_name__							:		str       ,       __magic_name__							:		Union[str, Any]       ,       __magic_name__							:		List[str]       ,       __magic_name__							:		Optional[Any]       ,       __magic_name__							:		Tuple       ,       **__magic_name__							:		Union[str, Any]       ,       ):







						"""simple docstring"""







						UpperCamelCase         =					True
						UpperCamelCase         =					BertGenerationEncoder(config=__magic_name__						)
						model.to(__magic_name__						)
						model.eval()
						UpperCamelCase         =					model(
						    __magic_name__       ,       attention_mask=__magic_name__       ,       encoder_hidden_states=__magic_name__       ,       encoder_attention_mask=__magic_name__       ,       )
						UpperCamelCase         =					model(
						    __magic_name__       ,       attention_mask=__magic_name__       ,       encoder_hidden_states=__magic_name__       ,       )
						self.parent.assertEqual(result.last_hidden_state.shape       ,       (self.batch_size, self.seq_length, self.hidden_size)						)




			def   lowerCamelCase_					(						self							:		Optional[int]       ,       __magic_name__							:		Optional[Any]       ,       __magic_name__							:		str       ,       __magic_name__							:		List[str]       ,       __magic_name__							:		Optional[int]       ,       __magic_name__							:		str       ,       __magic_name__							:		Union[str, Any]       ,       **__magic_name__							:		Dict       ,       ):







						"""simple docstring"""







						UpperCamelCase         =					True
						UpperCamelCase         =					True
						UpperCamelCase         =					BertGenerationDecoder(config=__magic_name__						).to(__magic_name__						).eval()

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

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

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

						UpperCamelCase         =					model(
						    __magic_name__       ,       attention_mask=__magic_name__       ,       encoder_hidden_states=__magic_name__       ,       encoder_attention_mask=__magic_name__       ,       output_hidden_states=__magic_name__       ,       )["""hidden_states"""][0]
						UpperCamelCase         =					model(
						    __magic_name__       ,       attention_mask=__magic_name__       ,       encoder_hidden_states=__magic_name__       ,       encoder_attention_mask=__magic_name__       ,       past_key_values=__magic_name__       ,       output_hidden_states=__magic_name__       ,       )["""hidden_states"""][0]

						# select random slice
						UpperCamelCase         =					ids_tensor((1,)       ,       output_from_past.shape[-1]						).item()
						UpperCamelCase         =					output_from_no_past[:, -3:, random_slice_idx].detach()
						UpperCamelCase         =					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(__magic_name__       ,       __magic_name__       ,       atol=1e-3						)						)




			def   lowerCamelCase_					(						self							:		int       ,       __magic_name__							:		Any       ,       __magic_name__							:		List[str]       ,       __magic_name__							:		Union[str, Any]       ,       __magic_name__							:		List[Any]       ,       *__magic_name__							:		Optional[Any]       ,       ):







						"""simple docstring"""







						UpperCamelCase         =					BertGenerationDecoder(__magic_name__						)
						model.to(__magic_name__						)
						model.eval()
						UpperCamelCase         =					model(__magic_name__       ,       attention_mask=__magic_name__       ,       labels=__magic_name__						)
						self.parent.assertEqual(result.logits.shape       ,       (self.batch_size, self.seq_length, self.vocab_size)						)





			def   lowerCamelCase_					(						self							:		List[Any]						):







						"""simple docstring"""







						UpperCamelCase  ,  UpperCamelCase  ,  UpperCamelCase  ,  UpperCamelCase         =					self.prepare_config_and_inputs()
						UpperCamelCase         =					{"""input_ids""": input_ids, """attention_mask""": input_mask}
						return config, inputs_dict




@require_torch
class    UpperCAmelCase					(       __snake_case      , __snake_case      , __snake_case      , unittest.TestCase    ):
			lowercase						= (BertGenerationEncoder, BertGenerationDecoder) if is_torch_available() else ()
			lowercase						= (BertGenerationDecoder,) if is_torch_available() else ()
			lowercase						= (
			    {"""feature-extraction""": BertGenerationEncoder, """text-generation""": BertGenerationDecoder}
			    if is_torch_available()
			    else {}
			)




			def   lowerCamelCase_					(						self							:		Tuple						):







						"""simple docstring"""







						UpperCamelCase         =					BertGenerationEncoderTester(self						)
						UpperCamelCase         =					ConfigTester(self       ,       config_class=__magic_name__       ,       hidden_size=3_7						)




			def   lowerCamelCase_					(						self							:		Optional[int]						):







						"""simple docstring"""







						self.config_tester.run_common_tests()




			def   lowerCamelCase_					(						self							:		Any						):







						"""simple docstring"""







						UpperCamelCase         =					self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*__magic_name__						)




			def   lowerCamelCase_					(						self							:		List[str]						):







						"""simple docstring"""







						UpperCamelCase  ,  UpperCamelCase  ,  UpperCamelCase  ,  UpperCamelCase         =					self.model_tester.prepare_config_and_inputs()
						UpperCamelCase         =					"""bert"""
						self.model_tester.create_and_check_model(__magic_name__       ,       __magic_name__       ,       __magic_name__       ,       __magic_name__						)




			def   lowerCamelCase_					(						self							:		Tuple						):







						"""simple docstring"""







						UpperCamelCase         =					self.model_tester.prepare_config_and_inputs_for_decoder()
						self.model_tester.create_and_check_model_as_decoder(*__magic_name__						)




			def   lowerCamelCase_					(						self							:		List[Any]						):







						"""simple docstring"""







						UpperCamelCase         =					self.model_tester.prepare_config_and_inputs_for_decoder()
						self.model_tester.create_and_check_decoder_model_past_large_inputs(*__magic_name__						)




			def   lowerCamelCase_					(						self							:		Dict						):







						"""simple docstring"""







						(
						    (
						    UpperCamelCase
						)  ,  (
						    UpperCamelCase
						)  ,  (
						    UpperCamelCase
						)  ,  (
						    UpperCamelCase
						)  ,  (
						    UpperCamelCase
						)  ,  (
						    UpperCamelCase
						)  ,  
						)         =					self.model_tester.prepare_config_and_inputs_for_decoder()

						UpperCamelCase         =					None

						self.model_tester.create_and_check_model_as_decoder(
						    __magic_name__       ,       __magic_name__       ,       __magic_name__       ,       __magic_name__       ,       __magic_name__       ,       __magic_name__       ,       )




			def   lowerCamelCase_					(						self							:		str						):







						"""simple docstring"""







						UpperCamelCase         =					self.model_tester.prepare_config_and_inputs_for_decoder()
						self.model_tester.create_and_check_for_causal_lm(*__magic_name__						)





			@slow
			def   lowerCamelCase_					(						self							:		List[Any]						):







						"""simple docstring"""







						UpperCamelCase         =					BertGenerationEncoder.from_pretrained("""google/bert_for_seq_generation_L-24_bbc_encoder"""						)
						self.assertIsNotNone(__magic_name__						)




@require_torch
class    UpperCAmelCase					(       unittest.TestCase    ):




			@slow
			def   lowerCamelCase_					(						self							:		Tuple						):







						"""simple docstring"""







						UpperCamelCase         =					BertGenerationEncoder.from_pretrained("""google/bert_for_seq_generation_L-24_bbc_encoder"""						)
						UpperCamelCase         =					torch.tensor([[1_0_1, 7_5_9_2, 1_0_1_0, 2_0_2_6, 3_8_9_9, 2_0_0_3, 1_0_1_4_0, 1_0_2]]						)
						with torch.no_grad():
									UpperCamelCase         =					model(__magic_name__						)[0]
						UpperCamelCase         =					torch.Size([1, 8, 1_0_2_4]						)
						self.assertEqual(output.shape       ,       __magic_name__						)
						UpperCamelCase         =					torch.tensor(
						    [[[0.1_775, 0.0_083, -0.0_321], [1.6_002, 0.1_287, 0.3_912], [2.1_473, 0.5_791, 0.6_066]]]						)
						self.assertTrue(torch.allclose(output[:, :3, :3]       ,       __magic_name__       ,       atol=1e-4						)						)




@require_torch
class    UpperCAmelCase					(       unittest.TestCase    ):




			@slow
			def   lowerCamelCase_					(						self							:		int						):







						"""simple docstring"""







						UpperCamelCase         =					BertGenerationDecoder.from_pretrained("""google/bert_for_seq_generation_L-24_bbc_encoder"""						)
						UpperCamelCase         =					torch.tensor([[1_0_1, 7_5_9_2, 1_0_1_0, 2_0_2_6, 3_8_9_9, 2_0_0_3, 1_0_1_4_0, 1_0_2]]						)
						with torch.no_grad():
									UpperCamelCase         =					model(__magic_name__						)[0]
						UpperCamelCase         =					torch.Size([1, 8, 5_0_3_5_8]						)
						self.assertEqual(output.shape       ,       __magic_name__						)
						UpperCamelCase         =					torch.tensor(
						    [[[-0.5_788, -2.5_994, -3.7_054], [0.0_438, 4.7_997, 1.8_795], [1.5_862, 6.6_409, 4.4_638]]]						)
						self.assertTrue(torch.allclose(output[:, :3, :3]       ,       __magic_name__       ,       atol=1e-4						)						)




