


from collections.abc import Generator
from math import sin






def 	__lowerCamelCase						(							lowerCamelCase__		:						bytes      ):





			'''simple docstring'''




			if len(lowerCamelCase__      ) != 32:
						raise ValueError("""Input must be of length 32"""      )

			lowerCamelCase    		=  b""""""
			for i in [3, 2, 1, 0]:
						little_endian += string_aa[8 * i : 8 * i + 8]
			return little_endian






def 	__lowerCamelCase						(							lowerCamelCase__		:						int      ):





			'''simple docstring'''




			if i < 0:
						raise ValueError("""Input must be non-negative"""      )

			lowerCamelCase    		=  format(lowerCamelCase__  ,							"""08x"""      )[-8:]
			lowerCamelCase    		=  b""""""
			for i in [3, 2, 1, 0]:
						little_endian_hex += hex_rep[2 * i : 2 * i + 2].encode("""utf-8"""      )
			return little_endian_hex






def 	__lowerCamelCase						(							lowerCamelCase__		:						bytes      ):





			'''simple docstring'''




			lowerCamelCase    		=  b""""""
			for char in message:
						bit_string += format(lowerCamelCase__  ,							"""08b"""      ).encode("""utf-8"""      )
			lowerCamelCase    		=  format(len(lowerCamelCase__      )  ,							"""064b"""      ).encode("""utf-8"""      )

			# Pad bit_string to a multiple of 512 chars
			bit_string += b"1"
			while len(lowerCamelCase__      ) % 512 != 448:
						bit_string += b"0"
			bit_string += to_little_endian(start_len[32:]      ) + to_little_endian(start_len[:32]      )

			return bit_string






def 	__lowerCamelCase						(							lowerCamelCase__		:						bytes      ):





			'''simple docstring'''




			if len(lowerCamelCase__      ) % 512 != 0:
						raise ValueError("""Input must have length that's a multiple of 512"""      )

			for pos in range(0  ,							len(lowerCamelCase__      )  ,							512      ):
						lowerCamelCase    		=  bit_string[pos : pos + 512]
						lowerCamelCase    		=  []
						for i in range(0  ,							512  ,							32      ):
									block_words.append(int(to_little_endian(block[i : i + 32]      )  ,							2      )      )
						yield block_words






def 	__lowerCamelCase						(							lowerCamelCase__		:						int      ):





			'''simple docstring'''




			if i < 0:
						raise ValueError("""Input must be non-negative"""      )

			lowerCamelCase    		=  format(lowerCamelCase__  ,							"""032b"""      )
			lowerCamelCase    		=  """"""
			for c in i_str:
						new_str += "1" if c == "0" else "0"
			return int(lowerCamelCase__  ,							2      )






def 	__lowerCamelCase						(							lowerCamelCase__		:						int  ,							lowerCamelCase__		:						int      ):





			'''simple docstring'''




			return (a + b) % 2**32






def 	__lowerCamelCase						(							lowerCamelCase__		:						int  ,							lowerCamelCase__		:						int      ):





			'''simple docstring'''




			if i < 0:
						raise ValueError("""Input must be non-negative"""      )
			if shift < 0:
						raise ValueError("""Shift must be non-negative"""      )
			return ((i << shift) ^ (i >> (32 - shift))) % 2**32






def 	__lowerCamelCase						(							lowerCamelCase__		:						bytes      ):





			'''simple docstring'''




			lowerCamelCase    		=  preprocess(lowerCamelCase__      )

			lowerCamelCase    		=  [int(2**32 * abs(sin(i + 1      )      )      ) for i in range(64      )]

			# Starting states
			lowerCamelCase    		=  0X6745_2301
			lowerCamelCase    		=  0XEFCD_AB89
			lowerCamelCase    		=  0X98BA_DCFE
			lowerCamelCase    		=  0X1032_5476

			lowerCamelCase    		=  [
			    7,
			    12,
			    17,
			    22,
			    7,
			    12,
			    17,
			    22,
			    7,
			    12,
			    17,
			    22,
			    7,
			    12,
			    17,
			    22,
			    5,
			    9,
			    14,
			    20,
			    5,
			    9,
			    14,
			    20,
			    5,
			    9,
			    14,
			    20,
			    5,
			    9,
			    14,
			    20,
			    4,
			    11,
			    16,
			    23,
			    4,
			    11,
			    16,
			    23,
			    4,
			    11,
			    16,
			    23,
			    4,
			    11,
			    16,
			    23,
			    6,
			    10,
			    15,
			    21,
			    6,
			    10,
			    15,
			    21,
			    6,
			    10,
			    15,
			    21,
			    6,
			    10,
			    15,
			    21,
			]

			# Process bit string in chunks, each with 16 32-char words
			for block_words in get_block_words(lowerCamelCase__      ):
						lowerCamelCase    		=  aa
						lowerCamelCase    		=  ba
						lowerCamelCase    		=  ca
						lowerCamelCase    		=  da

						# Hash current chunk
						for i in range(64      ):
									if i <= 15:
												# f = (b & c) | (not_32(b) & d)     # Alternate definition for f
												lowerCamelCase    		=  d ^ (b & (c ^ d))
												lowerCamelCase    		=  i
									elif i <= 31:
												# f = (d & b) | (not_32(d) & c)     # Alternate definition for f
												lowerCamelCase    		=  c ^ (d & (b ^ c))
												lowerCamelCase    		=  (5 * i + 1) % 16
									elif i <= 47:
												lowerCamelCase    		=  b ^ c ^ d
												lowerCamelCase    		=  (3 * i + 5) % 16
									else:
												lowerCamelCase    		=  c ^ (b | not_aa(lowerCamelCase__      ))
												lowerCamelCase    		=  (7 * i) % 16
									lowerCamelCase    		=  (f + a + added_consts[i] + block_words[g]) % 2**32
									lowerCamelCase    		=  d
									lowerCamelCase    		=  c
									lowerCamelCase    		=  b
									lowerCamelCase    		=  sum_aa(lowerCamelCase__  ,							left_rotate_aa(lowerCamelCase__  ,							shift_amounts[i]      )      )

						# Add hashed chunk to running total
						lowerCamelCase    		=  sum_aa(lowerCamelCase__  ,							lowerCamelCase__      )
						lowerCamelCase    		=  sum_aa(lowerCamelCase__  ,							lowerCamelCase__      )
						lowerCamelCase    		=  sum_aa(lowerCamelCase__  ,							lowerCamelCase__      )
						lowerCamelCase    		=  sum_aa(lowerCamelCase__  ,							lowerCamelCase__      )

			lowerCamelCase    		=  reformat_hex(lowerCamelCase__      ) + reformat_hex(lowerCamelCase__      ) + reformat_hex(lowerCamelCase__      ) + reformat_hex(lowerCamelCase__      )
			return digest


if __name__ == "__main__":
		import doctest

		doctest.testmod()





import argparse
import logging
import os

import datasets
import tensorflow as tf

from transformers import AutoTokenizer


UpperCAmelCase   :    Any     =		logging.getLogger(__name__)






def 	__lowerCamelCase						(							):





			'''simple docstring'''




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

			lowerCamelCase    		=  parser.parse_args()
			return args






def 	__lowerCamelCase						(							lowerCamelCase__		:						Any      ):





			'''simple docstring'''




			def fn(lowerCamelCase__		:						Any      ):
						return tokenizer(examples["""text"""]      )

			return fn






def 	__lowerCamelCase						(							lowerCamelCase__		:						Optional[int]      ):





			'''simple docstring'''




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






def 	__lowerCamelCase						(							lowerCamelCase__		:						str      ):





			'''simple docstring'''




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

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

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

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

			lowerCamelCase    		=  dataset_tokenized.map(lowerCamelCase__  ,							batched=lowerCamelCase__  ,							batch_size=1000  ,							num_proc=4      )

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

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

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


if __name__ == "__main__":
		UpperCAmelCase   :    Optional[int]     =		parse_args()
		main(args)

