
"""simple docstring"""


import itertools
import json
import os
import unittest

from transformers import AddedToken, LongformerTokenizer, LongformerTokenizerFast
from transformers.models.longformer.tokenization_longformer import VOCAB_FILES_NAMES
from transformers.testing_utils import require_tokenizers, slow

from ...test_tokenization_common import TokenizerTesterMixin



@require_tokenizers
class 	__a	(   _lowerCAmelCase     ,  unittest.TestCase						):
       UpperCamelCase_		:       Optional[int]				=      LongformerTokenizer
       UpperCamelCase_		:       str				=      True
       UpperCamelCase_		:       Tuple				=      LongformerTokenizerFast
       UpperCamelCase_		:       List[str]				=      True






       def 						_SCREAMING_SNAKE_CASE							(	self			:				Dict      )->	Any:







              """simple docstring"""
              super().setUp()

              # Adapted from Sennrich et al. 2015 and https://github.com/rsennrich/subword-nmt
              UpperCamelCase       			=		[
                  "l",
                  "o",
                  "w",
                  "e",
                  "r",
                  "s",
                  "t",
                  "i",
                  "d",
                  "n",
                  "\u0120",
                  "\u0120l",
                  "\u0120n",
                  "\u0120lo",
                  "\u0120low",
                  "er",
                  "\u0120lowest",
                  "\u0120newer",
                  "\u0120wider",
                  "<unk>",
              ]
              UpperCamelCase       			=		dict(zip(UpperCAmelCase_							,   range(len(UpperCAmelCase_      )      )      )      )
              UpperCamelCase       			=		["#version: 0.2", "\u0120 l", "\u0120l o", "\u0120lo w", "e r", ""]
              UpperCamelCase       			=		{"unk_token": "<unk>"}

              UpperCamelCase       			=		os.path.join(self.tmpdirname							,   VOCAB_FILES_NAMES["vocab_file"]      )
              UpperCamelCase       			=		os.path.join(self.tmpdirname							,   VOCAB_FILES_NAMES["merges_file"]      )
              with open(self.vocab_file							,   "w"							,   encoding="utf-8"      ) as fp:
                     fp.write(json.dumps(UpperCAmelCase_      ) + "\n"      )
              with open(self.merges_file							,   "w"							,   encoding="utf-8"      ) as fp:
                     fp.write("\n".join(UpperCAmelCase_      )      )






       def 						_SCREAMING_SNAKE_CASE							(	self			:				List[Any]							,   **UpperCAmelCase_			:				Dict      )->	int:







              """simple docstring"""
              kwargs.update(self.special_tokens_map      )
              return self.tokenizer_class.from_pretrained(self.tmpdirname							,   **UpperCAmelCase_      )






       def 						_SCREAMING_SNAKE_CASE							(	self			:				str							,   **UpperCAmelCase_			:				Optional[Any]      )->	Optional[int]:







              """simple docstring"""
              kwargs.update(self.special_tokens_map      )
              return self.rust_tokenizer_class.from_pretrained(self.tmpdirname							,   **UpperCAmelCase_      )






       def 						_SCREAMING_SNAKE_CASE							(	self			:				Optional[int]							,   UpperCAmelCase_			:				Dict      )->	List[Any]:







              """simple docstring"""
              UpperCamelCase       			=		"lower newer"
              UpperCamelCase       			=		"lower newer"
              return input_text, output_text






       def 						_SCREAMING_SNAKE_CASE							(	self			:				Tuple      )->	Optional[Any]:







              """simple docstring"""
              UpperCamelCase       			=		self.tokenizer_class(self.vocab_file							,   self.merges_file							,   **self.special_tokens_map      )
              UpperCamelCase       			=		"lower newer"
              UpperCamelCase       			=		["l", "o", "w", "er", "\u0120", "n", "e", "w", "er"]
              UpperCamelCase       			=		tokenizer.tokenize(UpperCAmelCase_      )  # , add_prefix_space=True)
              self.assertListEqual(UpperCAmelCase_							,   UpperCAmelCase_      )

              UpperCamelCase       			=		tokens + [tokenizer.unk_token]
              UpperCamelCase       			=		[0, 1, 2, 15, 10, 9, 3, 2, 15, 19]
              self.assertListEqual(tokenizer.convert_tokens_to_ids(UpperCAmelCase_      )							,   UpperCAmelCase_      )






       def 						_SCREAMING_SNAKE_CASE							(	self			:				Optional[int]      )->	Tuple:







              """simple docstring"""
              UpperCamelCase       			=		self.get_tokenizer()

              self.assertListEqual(tokenizer.encode("Hello world!"							,   add_special_tokens=UpperCAmelCase_      )							,   [0, 31_414, 232, 328, 2]      )
              self.assertListEqual(
                  tokenizer.encode("Hello world! cécé herlolip 418"							,   add_special_tokens=UpperCAmelCase_      )							,   [0, 31_414, 232, 328, 740, 1_140, 12_695, 69, 46_078, 1_588, 2]							,   )






       @slow
       def 						_SCREAMING_SNAKE_CASE							(	self			:				int      )->	int:







              """simple docstring"""
              UpperCamelCase       			=		self.tokenizer_class.from_pretrained("allenai/longformer-base-4096"      )

              UpperCamelCase       			=		tokenizer.encode("sequence builders"							,   add_special_tokens=UpperCAmelCase_      )
              UpperCamelCase       			=		tokenizer.encode("multi-sequence build"							,   add_special_tokens=UpperCAmelCase_      )

              UpperCamelCase       			=		tokenizer.encode(
                  "sequence builders"							,   add_special_tokens=UpperCAmelCase_							,   add_prefix_space=UpperCAmelCase_      )
              UpperCamelCase       			=		tokenizer.encode(
                  "sequence builders"							,   "multi-sequence build"							,   add_special_tokens=UpperCAmelCase_							,   add_prefix_space=UpperCAmelCase_      )

              UpperCamelCase       			=		tokenizer.build_inputs_with_special_tokens(UpperCAmelCase_      )
              UpperCamelCase       			=		tokenizer.build_inputs_with_special_tokens(UpperCAmelCase_							,   UpperCAmelCase_      )

              assert encoded_sentence == encoded_text_from_decode
              assert encoded_pair == encoded_pair_from_decode






       def 						_SCREAMING_SNAKE_CASE							(	self			:				Optional[int]      )->	Optional[Any]:







              """simple docstring"""
              UpperCamelCase       			=		self.get_tokenizer()

              UpperCamelCase       			=		"Encode this sequence."
              UpperCamelCase       			=		tokenizer.byte_encoder[" ".encode("utf-8"      )[0]]

              # Testing encoder arguments
              UpperCamelCase       			=		tokenizer.encode(UpperCAmelCase_							,   add_special_tokens=UpperCAmelCase_							,   add_prefix_space=UpperCAmelCase_      )
              UpperCamelCase       			=		tokenizer.convert_ids_to_tokens(encoded[0]      )[0]
              self.assertNotEqual(UpperCAmelCase_							,   UpperCAmelCase_      )

              UpperCamelCase       			=		tokenizer.encode(UpperCAmelCase_							,   add_special_tokens=UpperCAmelCase_							,   add_prefix_space=UpperCAmelCase_      )
              UpperCamelCase       			=		tokenizer.convert_ids_to_tokens(encoded[0]      )[0]
              self.assertEqual(UpperCAmelCase_							,   UpperCAmelCase_      )

              tokenizer.add_special_tokens({"bos_token": "<s>"}      )
              UpperCamelCase       			=		tokenizer.encode(UpperCAmelCase_							,   add_special_tokens=UpperCAmelCase_      )
              UpperCamelCase       			=		tokenizer.convert_ids_to_tokens(encoded[1]      )[0]
              self.assertNotEqual(UpperCAmelCase_							,   UpperCAmelCase_      )

              # Testing spaces after special tokens
              UpperCamelCase       			=		"<mask>"
              tokenizer.add_special_tokens(
                  {"mask_token": AddedToken(UpperCAmelCase_							,   lstrip=UpperCAmelCase_							,   rstrip=UpperCAmelCase_      )}      )  # mask token has a left space
              UpperCamelCase       			=		tokenizer.convert_tokens_to_ids(UpperCAmelCase_      )

              UpperCamelCase       			=		"Encode <mask> sequence"
              UpperCamelCase       			=		"Encode <mask>sequence"

              UpperCamelCase       			=		tokenizer.encode(UpperCAmelCase_      )
              UpperCamelCase       			=		encoded.index(UpperCAmelCase_      )
              UpperCamelCase       			=		tokenizer.convert_ids_to_tokens(encoded[mask_loc + 1]      )[0]
              self.assertEqual(UpperCAmelCase_							,   UpperCAmelCase_      )

              UpperCamelCase       			=		tokenizer.encode(UpperCAmelCase_      )
              UpperCamelCase       			=		encoded.index(UpperCAmelCase_      )
              UpperCamelCase       			=		tokenizer.convert_ids_to_tokens(encoded[mask_loc + 1]      )[0]
              self.assertNotEqual(UpperCAmelCase_							,   UpperCAmelCase_      )






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







              """simple docstring"""
              pass






       def 						_SCREAMING_SNAKE_CASE							(	self			:				Optional[Any]      )->	Tuple:







              """simple docstring"""
              for tokenizer, pretrained_name, kwargs in self.tokenizers_list:
                     with self.subTest(f"{tokenizer.__class__.__name__} ({pretrained_name})"      ):
                            UpperCamelCase       			=		self.rust_tokenizer_class.from_pretrained(UpperCAmelCase_							,   **UpperCAmelCase_      )
                            UpperCamelCase       			=		self.tokenizer_class.from_pretrained(UpperCAmelCase_							,   **UpperCAmelCase_      )
                            UpperCamelCase       			=		"A, <mask> AllenNLP sentence."
                            UpperCamelCase       			=		tokenizer_r.encode_plus(UpperCAmelCase_							,   add_special_tokens=UpperCAmelCase_							,   return_token_type_ids=UpperCAmelCase_      )
                            UpperCamelCase       			=		tokenizer_p.encode_plus(UpperCAmelCase_							,   add_special_tokens=UpperCAmelCase_							,   return_token_type_ids=UpperCAmelCase_      )

                            # token_type_ids should put 0 everywhere
                            self.assertEqual(sum(tokens_r["token_type_ids"]      )							,   sum(tokens_p["token_type_ids"]      )      )

                            # attention_mask should put 1 everywhere, so sum over length should be 1
                            self.assertEqual(
                                sum(tokens_r["attention_mask"]      ) / len(tokens_r["attention_mask"]      )							,   sum(tokens_p["attention_mask"]      ) / len(tokens_p["attention_mask"]      )							,   )

                            UpperCamelCase       			=		tokenizer_r.convert_ids_to_tokens(tokens_r["input_ids"]      )
                            UpperCamelCase       			=		tokenizer_p.convert_ids_to_tokens(tokens_p["input_ids"]      )

                            # Rust correctly handles the space before the mask while python doesnt
                            self.assertSequenceEqual(tokens_p["input_ids"]							,   [0, 250, 6, 50_264, 3_823, 487, 21_992, 3_645, 4, 2]      )
                            self.assertSequenceEqual(tokens_r["input_ids"]							,   [0, 250, 6, 50_264, 3_823, 487, 21_992, 3_645, 4, 2]      )

                            self.assertSequenceEqual(
                                UpperCAmelCase_							,   ["<s>", "A", ",", "<mask>", "ĠAllen", "N", "LP", "Ġsentence", ".", "</s>"]      )
                            self.assertSequenceEqual(
                                UpperCAmelCase_							,   ["<s>", "A", ",", "<mask>", "ĠAllen", "N", "LP", "Ġsentence", ".", "</s>"]      )






       def 						_SCREAMING_SNAKE_CASE							(	self			:				Any      )->	Dict:







              """simple docstring"""
              for trim_offsets, add_prefix_space in itertools.product([True, False]							,   repeat=2      ):
                     UpperCamelCase       			=		self.rust_tokenizer_class.from_pretrained(
                         self.tmpdirname							,   use_fast=UpperCAmelCase_							,   add_prefix_space=UpperCAmelCase_							,   trim_offsets=UpperCAmelCase_      )

                     UpperCamelCase       			=		json.loads(tokenizer_r.backend_tokenizer.pre_tokenizer.__getstate__()      )
                     UpperCamelCase       			=		json.loads(tokenizer_r.backend_tokenizer.post_processor.__getstate__()      )

                     self.assertEqual(pre_tokenizer_state["add_prefix_space"]							,   UpperCAmelCase_      )

                     self.assertEqual(post_processor_state["add_prefix_space"]							,   UpperCAmelCase_      )
                     self.assertEqual(post_processor_state["trim_offsets"]							,   UpperCAmelCase_      )






       def 						_SCREAMING_SNAKE_CASE							(	self			:				Optional[int]      )->	Union[str, Any]:







              """simple docstring"""
              # Test which aims to verify that the offsets are well adapted to the argument `add_prefix_space` and
              # `trim_offsets`
              for tokenizer, pretrained_name, kwargs in self.tokenizers_list:
                     with self.subTest(f"{tokenizer.__class__.__name__} ({pretrained_name})"      ):
                            UpperCamelCase       			=		"hello"  # `hello` is a token in the vocabulary of `pretrained_name`
                            UpperCamelCase       			=		f"{text_of_1_token} {text_of_1_token}"

                            UpperCamelCase       			=		self.rust_tokenizer_class.from_pretrained(
                                UpperCAmelCase_							,   use_fast=UpperCAmelCase_							,   add_prefix_space=UpperCAmelCase_							,   trim_offsets=UpperCAmelCase_      )
                            UpperCamelCase       			=		tokenizer_r(UpperCAmelCase_							,   return_offsets_mapping=UpperCAmelCase_							,   add_special_tokens=UpperCAmelCase_      )
                            self.assertEqual(encoding.offset_mapping[0]							,   (0, len(UpperCAmelCase_      ))      )
                            self.assertEqual(
                                encoding.offset_mapping[1]							,   (len(UpperCAmelCase_      ) + 1, len(UpperCAmelCase_      ) + 1 + len(UpperCAmelCase_      ))							,   )

                            UpperCamelCase       			=		self.rust_tokenizer_class.from_pretrained(
                                UpperCAmelCase_							,   use_fast=UpperCAmelCase_							,   add_prefix_space=UpperCAmelCase_							,   trim_offsets=UpperCAmelCase_      )
                            UpperCamelCase       			=		tokenizer_r(UpperCAmelCase_							,   return_offsets_mapping=UpperCAmelCase_							,   add_special_tokens=UpperCAmelCase_      )
                            self.assertEqual(encoding.offset_mapping[0]							,   (0, len(UpperCAmelCase_      ))      )
                            self.assertEqual(
                                encoding.offset_mapping[1]							,   (len(UpperCAmelCase_      ) + 1, len(UpperCAmelCase_      ) + 1 + len(UpperCAmelCase_      ))							,   )

                            UpperCamelCase       			=		self.rust_tokenizer_class.from_pretrained(
                                UpperCAmelCase_							,   use_fast=UpperCAmelCase_							,   add_prefix_space=UpperCAmelCase_							,   trim_offsets=UpperCAmelCase_      )
                            UpperCamelCase       			=		tokenizer_r(UpperCAmelCase_							,   return_offsets_mapping=UpperCAmelCase_							,   add_special_tokens=UpperCAmelCase_      )
                            self.assertEqual(encoding.offset_mapping[0]							,   (0, len(UpperCAmelCase_      ))      )
                            self.assertEqual(
                                encoding.offset_mapping[1]							,   (len(UpperCAmelCase_      ), len(UpperCAmelCase_      ) + 1 + len(UpperCAmelCase_      ))							,   )

                            UpperCamelCase       			=		self.rust_tokenizer_class.from_pretrained(
                                UpperCAmelCase_							,   use_fast=UpperCAmelCase_							,   add_prefix_space=UpperCAmelCase_							,   trim_offsets=UpperCAmelCase_      )
                            UpperCamelCase       			=		tokenizer_r(UpperCAmelCase_							,   return_offsets_mapping=UpperCAmelCase_							,   add_special_tokens=UpperCAmelCase_      )
                            self.assertEqual(encoding.offset_mapping[0]							,   (0, len(UpperCAmelCase_      ))      )
                            self.assertEqual(
                                encoding.offset_mapping[1]							,   (len(UpperCAmelCase_      ), len(UpperCAmelCase_      ) + 1 + len(UpperCAmelCase_      ))							,   )

                            UpperCamelCase       			=		f" {text}"

                            # tokenizer_r = self.rust_tokenizer_class.from_pretrained(
                            #     pretrained_name, use_fast=True, add_prefix_space=True, trim_offsets=True
                            # )
                            # encoding = tokenizer_r(text, return_offsets_mapping=True, add_special_tokens=False)
                            # self.assertEqual(encoding.offset_mapping[0], (1, 1 + len(text_of_1_token)))
                            # self.assertEqual(
                            #     encoding.offset_mapping[1],
                            #     (1 + len(text_of_1_token) + 1, 1 + len(text_of_1_token) + 1 + len(text_of_1_token)),
                            # )

                            UpperCamelCase       			=		self.rust_tokenizer_class.from_pretrained(
                                UpperCAmelCase_							,   use_fast=UpperCAmelCase_							,   add_prefix_space=UpperCAmelCase_							,   trim_offsets=UpperCAmelCase_      )
                            UpperCamelCase       			=		tokenizer_r(UpperCAmelCase_							,   return_offsets_mapping=UpperCAmelCase_							,   add_special_tokens=UpperCAmelCase_      )
                            self.assertEqual(encoding.offset_mapping[0]							,   (1, 1 + len(UpperCAmelCase_      ))      )
                            self.assertEqual(
                                encoding.offset_mapping[1]							,   (1 + len(UpperCAmelCase_      ) + 1, 1 + len(UpperCAmelCase_      ) + 1 + len(UpperCAmelCase_      ))							,   )

                            UpperCamelCase       			=		self.rust_tokenizer_class.from_pretrained(
                                UpperCAmelCase_							,   use_fast=UpperCAmelCase_							,   add_prefix_space=UpperCAmelCase_							,   trim_offsets=UpperCAmelCase_      )
                            UpperCamelCase       			=		tokenizer_r(UpperCAmelCase_							,   return_offsets_mapping=UpperCAmelCase_							,   add_special_tokens=UpperCAmelCase_      )
                            self.assertEqual(encoding.offset_mapping[0]							,   (0, 1 + len(UpperCAmelCase_      ))      )
                            self.assertEqual(
                                encoding.offset_mapping[1]							,   (1 + len(UpperCAmelCase_      ), 1 + len(UpperCAmelCase_      ) + 1 + len(UpperCAmelCase_      ))							,   )

                            UpperCamelCase       			=		self.rust_tokenizer_class.from_pretrained(
                                UpperCAmelCase_							,   use_fast=UpperCAmelCase_							,   add_prefix_space=UpperCAmelCase_							,   trim_offsets=UpperCAmelCase_      )
                            UpperCamelCase       			=		tokenizer_r(UpperCAmelCase_							,   return_offsets_mapping=UpperCAmelCase_							,   add_special_tokens=UpperCAmelCase_      )
                            self.assertEqual(encoding.offset_mapping[0]							,   (0, 1 + len(UpperCAmelCase_      ))      )
                            self.assertEqual(
                                encoding.offset_mapping[1]							,   (1 + len(UpperCAmelCase_      ), 1 + len(UpperCAmelCase_      ) + 1 + len(UpperCAmelCase_      ))							,   )







"""simple docstring"""


import argparse
import logging
import os

import datasets
import tensorflow as tf

from transformers import AutoTokenizer


SCREAMING_SNAKE_CASE			  =  logging.getLogger(__name__)



def 						lowerCamelCase__      (				)->				Tuple:







       """simple docstring"""
       UpperCamelCase       			=		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=10_00		,   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=5_12		,   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."		,   )

       UpperCamelCase       			=		parser.parse_args()
       return args



def 						lowerCamelCase__      (				UpperCAmelCase_					)->				Optional[int]:







       """simple docstring"""
       def fn(UpperCAmelCase_					):
              return tokenizer(examples["text"]					)

       return fn



def 						lowerCamelCase__      (				UpperCAmelCase_					)->				Union[str, Any]:







       """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=UpperCAmelCase_					)
              UpperCamelCase       			=		tf.train.Example(features=UpperCAmelCase_					)
              UpperCamelCase       			=		example.SerializeToString()
              records.append(UpperCAmelCase_					)
       return records



def 						lowerCamelCase__      (				UpperCAmelCase_					)->				int:







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

       if args.limit is not None:
              UpperCamelCase       			=		min(len(UpperCAmelCase_					)		,   args.limit					)
              UpperCamelCase       			=		dataset.select(range(UpperCAmelCase_					)					)
              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(UpperCAmelCase_					):
                     os.makedirs(UpperCAmelCase_					)
       else:
              UpperCamelCase       			=		os.path.join(args.output_dir		,   args.split					)

       # Tokenize the whole dataset at once.
       UpperCamelCase       			=		tokenize_function(UpperCAmelCase_					)
       UpperCamelCase       			=		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.
              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		,   UpperCAmelCase_		,   args.max_length					)]
                  for k, t in concatenated_examples.items()
              }
              return result

       UpperCamelCase       			=		dataset_tokenized.map(UpperCAmelCase_		,   batched=UpperCAmelCase_		,   batch_size=10_00		,   num_proc=4					)

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

              with tf.io.TFRecordWriter(UpperCAmelCase_					) as out_file:
                     for i in range(len(UpperCAmelCase_					)					):
                            UpperCamelCase       			=		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__":
      SCREAMING_SNAKE_CASE			  =  parse_args()
      main(args)





