





'''simple docstring'''





import json
from typing import Dict, List, Optional, Tuple, Union

from tokenizers import pre_tokenizers, processors

from ...tokenization_utils_base import AddedToken, BatchEncoding, EncodedInput
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import PaddingStrategy, logging
from .tokenization_led import LEDTokenizer


_lowerCamelCase   				=			logging.get_logger(__name__)


_lowerCamelCase   				=			{'vocab_file': 'vocab.json', 'merges_file': 'merges.txt', 'tokenizer_file': 'tokenizer.json'}

_lowerCamelCase   				=			{
    'vocab_file': {
        'allenai/led-base-16384': 'https://huggingface.co/allenai/led-base-16384/resolve/main/vocab.json',
    },
    'merges_file': {
        'allenai/led-base-16384': 'https://huggingface.co/allenai/led-base-16384/resolve/main/merges.txt',
    },
    'tokenizer_file': {
        'allenai/led-base-16384': 'https://huggingface.co/allenai/led-base-16384/resolve/main/tokenizer.json',
    },
}

_lowerCamelCase   				=			{
    'allenai/led-base-16384': 16384,
}




class   UpperCamelCase_       (     lowerCAmelCase__    ):
							lowerCamelCase_           =					VOCAB_FILES_NAMES
							lowerCamelCase_           =					PRETRAINED_VOCAB_FILES_MAP
							lowerCamelCase_           =					PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
							lowerCamelCase_           =					LEDTokenizer
							lowerCamelCase_           =					["input_ids", "attention_mask"]


							def __init__(							self    :Tuple							,				__A    :int=None							,				__A    :Dict=None							,				__A    :Union[str, Any]=None							,				__A    :List[str]="replace"							,				__A    :Optional[Any]="<s>"							,				__A    :str="</s>"							,				__A    :int="</s>"							,				__A    :List[str]="<s>"							,				__A    :Tuple="<unk>"							,				__A    :Optional[int]="<pad>"							,				__A    :Union[str, Any]="<mask>"							,				__A    :Union[str, Any]=False							,				__A    :int=True							,				**__A    :Dict							,				)   ->   Union[str, Any]:



									"""simple docstring"""



									super().__init__(
									    _SCREAMING_SNAKE_CASE							,				_SCREAMING_SNAKE_CASE							,				tokenizer_file=_SCREAMING_SNAKE_CASE							,				errors=_SCREAMING_SNAKE_CASE							,				bos_token=_SCREAMING_SNAKE_CASE							,				eos_token=_SCREAMING_SNAKE_CASE							,				sep_token=_SCREAMING_SNAKE_CASE							,				cls_token=_SCREAMING_SNAKE_CASE							,				unk_token=_SCREAMING_SNAKE_CASE							,				pad_token=_SCREAMING_SNAKE_CASE							,				mask_token=_SCREAMING_SNAKE_CASE							,				add_prefix_space=_SCREAMING_SNAKE_CASE							,				trim_offsets=_SCREAMING_SNAKE_CASE							,				**_SCREAMING_SNAKE_CASE							,				)

									SCREAMING_SNAKE_CASE__			=      json.loads(self.backend_tokenizer.pre_tokenizer.__getstate__()		)
									if pre_tok_state.get("""add_prefix_space"""							,				_SCREAMING_SNAKE_CASE		) != add_prefix_space:
											SCREAMING_SNAKE_CASE__			=      getattr(_SCREAMING_SNAKE_CASE							,				pre_tok_state.pop("""type"""		)		)
											SCREAMING_SNAKE_CASE__			=      add_prefix_space
											SCREAMING_SNAKE_CASE__			=      pre_tok_class(**_SCREAMING_SNAKE_CASE		)

									SCREAMING_SNAKE_CASE__			=      add_prefix_space

									# the pre_tokenizer is already updated in the GPT2TokenizerFast `__init__`
									SCREAMING_SNAKE_CASE__			=      """post_processor"""
									SCREAMING_SNAKE_CASE__			=      getattr(self.backend_tokenizer							,				_SCREAMING_SNAKE_CASE							,				_SCREAMING_SNAKE_CASE		)
									if tokenizer_component_instance:
											SCREAMING_SNAKE_CASE__			=      json.loads(tokenizer_component_instance.__getstate__()		)

											# The lists 'sep' and 'cls' must be cased in tuples for the object `post_processor_class`
											if "sep" in state:
													SCREAMING_SNAKE_CASE__			=      tuple(state["""sep"""]		)
											if "cls" in state:
													SCREAMING_SNAKE_CASE__			=      tuple(state["""cls"""]		)

											SCREAMING_SNAKE_CASE__			=      False

											if state.get("""add_prefix_space"""							,				_SCREAMING_SNAKE_CASE		) != add_prefix_space:
													SCREAMING_SNAKE_CASE__			=      add_prefix_space
													SCREAMING_SNAKE_CASE__			=      True

											if state.get("""trim_offsets"""							,				_SCREAMING_SNAKE_CASE		) != trim_offsets:
													SCREAMING_SNAKE_CASE__			=      trim_offsets
													SCREAMING_SNAKE_CASE__			=      True

											if changes_to_apply:
													SCREAMING_SNAKE_CASE__			=      getattr(_SCREAMING_SNAKE_CASE							,				state.pop("""type"""		)		)
													SCREAMING_SNAKE_CASE__			=      component_class(**_SCREAMING_SNAKE_CASE		)
													setattr(self.backend_tokenizer							,				_SCREAMING_SNAKE_CASE							,				_SCREAMING_SNAKE_CASE		)


							@property
							# Copied from transformers.models.bart.tokenization_bart_fast.BartTokenizerFast.mask_token with BART->LED
							def      _snake_case  (							self    :int		)   ->   str:



									"""simple docstring"""



									if self._mask_token is None:
											if self.verbose:
													logger.error("""Using mask_token, but it is not set yet."""		)
											return None
									return str(self._mask_token		)


							@mask_token.setter
							def      _snake_case  (							self    :List[Any]							,				__A    :int		)   ->   Dict:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      AddedToken(_SCREAMING_SNAKE_CASE							,				lstrip=_SCREAMING_SNAKE_CASE							,				rstrip=_SCREAMING_SNAKE_CASE		) if isinstance(_SCREAMING_SNAKE_CASE							,				_SCREAMING_SNAKE_CASE		) else value
									SCREAMING_SNAKE_CASE__			=      value


							def      _snake_case  (							self    :List[str]							,				*__A    :Optional[int]							,				**__A    :Optional[int]		)   ->   BatchEncoding:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      kwargs.get("""is_split_into_words"""							,				_SCREAMING_SNAKE_CASE		)

									if is_split_into_words and not self.add_prefix_space:
											raise ValueError(
											    f'''You need to instantiate {self.__class__.__name__} with add_prefix_space=True '''
											    """to use it with pretokenized inputs."""		)

									return super()._batch_encode_plus(*_SCREAMING_SNAKE_CASE							,				**_SCREAMING_SNAKE_CASE		)


							def      _snake_case  (							self    :Union[str, Any]							,				*__A    :List[Any]							,				**__A    :Union[str, Any]		)   ->   BatchEncoding:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      kwargs.get("""is_split_into_words"""							,				_SCREAMING_SNAKE_CASE		)

									if is_split_into_words and not self.add_prefix_space:
											raise ValueError(
											    f'''You need to instantiate {self.__class__.__name__} with add_prefix_space=True '''
											    """to use it with pretokenized inputs."""		)

									return super()._encode_plus(*_SCREAMING_SNAKE_CASE							,				**_SCREAMING_SNAKE_CASE		)


							def      _snake_case  (							self    :Union[str, Any]							,				__A    :Optional[int]							,				__A    :int = None		)   ->   Tuple[str]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      self._tokenizer.model.save(_SCREAMING_SNAKE_CASE							,				name=_SCREAMING_SNAKE_CASE		)
									return tuple(_SCREAMING_SNAKE_CASE		)


							def      _snake_case  (							self    :Union[str, Any]							,				__A    :Optional[int]							,				__A    :List[str]=None		)   ->   Optional[Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      [self.bos_token_id] + token_ids_a + [self.eos_token_id]
									if token_ids_a is None:
											return output

									return output + [self.eos_token_id] + token_ids_a + [self.eos_token_id]


							def      _snake_case  (							self    :Any							,				__A    :int							,				__A    :List[Any] = None		)   ->   List[int]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      [self.sep_token_id]
									SCREAMING_SNAKE_CASE__			=      [self.cls_token_id]

									if token_ids_a is None:
											return len(cls + token_ids_a + sep		) * [0]
									return len(cls + token_ids_a + sep + sep + token_ids_a + sep		) * [0]


							def      _snake_case  (							self    :Optional[int]							,				__A    :Union[str, Any]							,				__A    :Union[str, Any] = None							,				__A    :Optional[int] = PaddingStrategy.DO_NOT_PAD							,				__A    :Optional[Any] = None							,				__A    :Optional[Any] = None							,				)   ->   dict:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      super()._pad(
									    encoded_inputs=_SCREAMING_SNAKE_CASE							,				max_length=_SCREAMING_SNAKE_CASE							,				padding_strategy=_SCREAMING_SNAKE_CASE							,				pad_to_multiple_of=_SCREAMING_SNAKE_CASE							,				return_attention_mask=_SCREAMING_SNAKE_CASE							,				)

									# Load from model defaults
									if return_attention_mask is None:
											SCREAMING_SNAKE_CASE__			=      """attention_mask""" in self.model_input_names

									if return_attention_mask and "global_attention_mask" in encoded_inputs:
											SCREAMING_SNAKE_CASE__			=      encoded_inputs[self.model_input_names[0]]
											# `global_attention_mask` need to have the same length as other (sequential) inputs.
											SCREAMING_SNAKE_CASE__			=      len(encoded_inputs["""global_attention_mask"""]		) != len(_SCREAMING_SNAKE_CASE		)

											if needs_to_be_padded:
													SCREAMING_SNAKE_CASE__			=      len(_SCREAMING_SNAKE_CASE		) - len(encoded_inputs["""global_attention_mask"""]		)

													if self.padding_side == "right":
															# Use `-1` since `0` in `global_attention_mask` means `local attention` instead of `not to attend`
															SCREAMING_SNAKE_CASE__			=      (
															    encoded_inputs["""global_attention_mask"""] + [-1] * difference
															)
													elif self.padding_side == "left":
															SCREAMING_SNAKE_CASE__			=      [-1] * difference + encoded_inputs[
															    """global_attention_mask"""
															]
													else:
															raise ValueError("""Invalid padding strategy:""" + str(self.padding_side		)		)

									return encoded_inputs






import argparse
import logging
import os

import datasets
import tensorflow as tf

from transformers import AutoTokenizer


_lowerCamelCase   				=			logging.getLogger(__name__)






def 				SCREAMING_SNAKE_CASE__							(       ):
		SCREAMING_SNAKE_CASE__			=      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."""					,   )

		SCREAMING_SNAKE_CASE__			=      parser.parse_args()
		return args






def 				SCREAMING_SNAKE_CASE__							(       UpperCamelCase__:							List[Any]    ):
		def fn(UpperCamelCase__:							Any    ):
				return tokenizer(examples["""text"""]    )

		return fn






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






def 				SCREAMING_SNAKE_CASE__							(       UpperCamelCase__:							List[str]    ):
		SCREAMING_SNAKE_CASE__			=      datasets.load_dataset(args.dataset_name					,   args.dataset_config					,   split=args.split    )

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

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

		# Tokenize the whole dataset at once.
		SCREAMING_SNAKE_CASE__			=      tokenize_function(UpperCamelCase__    )
		SCREAMING_SNAKE_CASE__			=      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__:							int    ):
				# Concatenate all texts.
				SCREAMING_SNAKE_CASE__			=      {k: sum(examples[k]					,   []    ) for k in examples.keys()}
				SCREAMING_SNAKE_CASE__			=      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 🫀
				SCREAMING_SNAKE_CASE__			=      (total_length // args.max_length) * args.max_length
				# Split by chunks of max_len.
				SCREAMING_SNAKE_CASE__			=      {
				    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

		SCREAMING_SNAKE_CASE__			=      dataset_tokenized.map(UpperCamelCase__					,   batched=UpperCamelCase__					,   batch_size=1_000					,   num_proc=4    )

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

				with tf.io.TFRecordWriter(UpperCamelCase__    ) as out_file:
						for i in range(len(UpperCamelCase__    )    ):
								SCREAMING_SNAKE_CASE__			=      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)