



"""simple docstring"""







import inspect
import unittest

from transformers import DPTConfig
from transformers.file_utils import is_torch_available, is_vision_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_vision, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
     import torch
     from torch import nn

     from transformers import MODEL_MAPPING, DPTForDepthEstimation, DPTForSemanticSegmentation, DPTModel
     from transformers.models.dpt.modeling_dpt import DPT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
     from PIL import Image

     from transformers import DPTImageProcessor







class       UpperCamelCase__   :




      """simple docstring"""




      def __init__(      self				:      Dict    ,			SCREAMING_SNAKE_CASE_				:      List[str]    ,			SCREAMING_SNAKE_CASE_				:      List[Any]=2    ,			SCREAMING_SNAKE_CASE_				:      List[Any]=3_2    ,			SCREAMING_SNAKE_CASE_				:      Any=1_6    ,			SCREAMING_SNAKE_CASE_				:      Optional[Any]=3    ,			SCREAMING_SNAKE_CASE_				:      List[str]=True    ,			SCREAMING_SNAKE_CASE_				:      List[Any]=True    ,			SCREAMING_SNAKE_CASE_				:      Dict=3_2    ,			SCREAMING_SNAKE_CASE_				:      str=4    ,			SCREAMING_SNAKE_CASE_				:      List[str]=[0, 1, 2, 3]    ,			SCREAMING_SNAKE_CASE_				:      int=4    ,			SCREAMING_SNAKE_CASE_				:      int=3_7    ,			SCREAMING_SNAKE_CASE_				:      int="gelu"    ,			SCREAMING_SNAKE_CASE_				:      str=0.1    ,			SCREAMING_SNAKE_CASE_				:      str=0.1    ,			SCREAMING_SNAKE_CASE_				:      Optional[Any]=0.02    ,			SCREAMING_SNAKE_CASE_				:      Union[str, Any]=3    ,			SCREAMING_SNAKE_CASE_				:      Any=[1, 3_8_4, 2_4, 2_4]    ,			SCREAMING_SNAKE_CASE_				:      List[Any]=True    ,			SCREAMING_SNAKE_CASE_				:      Optional[Any]=None    ,			):
        lowerCAmelCase_   :				str								= parent
        lowerCAmelCase_   :				List[str]								= batch_size
        lowerCAmelCase_   :				int								= image_size
        lowerCAmelCase_   :				Any								= patch_size
        lowerCAmelCase_   :				List[str]								= num_channels
        lowerCAmelCase_   :				Union[str, Any]								= is_training
        lowerCAmelCase_   :				Union[str, Any]								= use_labels
        lowerCAmelCase_   :				Optional[int]								= hidden_size
        lowerCAmelCase_   :				List[str]								= num_hidden_layers
        lowerCAmelCase_   :				Tuple								= backbone_out_indices
        lowerCAmelCase_   :				Dict								= num_attention_heads
        lowerCAmelCase_   :				int								= intermediate_size
        lowerCAmelCase_   :				Optional[int]								= hidden_act
        lowerCAmelCase_   :				Dict								= hidden_dropout_prob
        lowerCAmelCase_   :				Optional[int]								= attention_probs_dropout_prob
        lowerCAmelCase_   :				List[Any]								= initializer_range
        lowerCAmelCase_   :				Tuple								= num_labels
        lowerCAmelCase_   :				Optional[int]								= backbone_featmap_shape
        lowerCAmelCase_   :				Tuple								= scope
        lowerCAmelCase_   :				int								= is_hybrid
        # sequence length of DPT = num_patches + 1 (we add 1 for the [CLS] token)
        lowerCAmelCase_   :				List[Any]								= (image_size // patch_size) ** 2
        lowerCAmelCase_   :				Union[str, Any]								= num_patches + 1



      def      SCREAMING_SNAKE_CASE__ (      self				:      List[str]       ):
        lowerCAmelCase_   :				List[str]								= floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]       )

        lowerCAmelCase_   :				Optional[Any]								= None
        if self.use_labels:
          lowerCAmelCase_   :				List[Any]								= ids_tensor([self.batch_size, self.image_size, self.image_size]    ,			self.num_labels       )

        lowerCAmelCase_   :				List[Any]								= self.get_config()

        return config, pixel_values, labels



      def      SCREAMING_SNAKE_CASE__ (      self				:      Dict       ):
        lowerCAmelCase_   :				List[Any]								= {
            'global_padding': 'same',
            'layer_type': 'bottleneck',
            'depths': [3, 4, 9],
            'out_features': ['stage1', 'stage2', 'stage3'],
            'embedding_dynamic_padding': True,
            'hidden_sizes': [9_6, 1_9_2, 3_8_4, 7_6_8],
            'num_groups': 2,
        }

        return DPTConfig(
            image_size=self.image_size    ,			patch_size=self.patch_size    ,			num_channels=self.num_channels    ,			hidden_size=self.hidden_size    ,			num_hidden_layers=self.num_hidden_layers    ,			backbone_out_indices=self.backbone_out_indices    ,			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    ,			is_decoder=SCREAMING_SNAKE_CASE_    ,			initializer_range=self.initializer_range    ,			is_hybrid=self.is_hybrid    ,			backbone_config=SCREAMING_SNAKE_CASE_    ,			backbone_featmap_shape=self.backbone_featmap_shape    ,			)



      def      SCREAMING_SNAKE_CASE__ (      self				:      Any    ,			SCREAMING_SNAKE_CASE_				:      Optional[int]    ,			SCREAMING_SNAKE_CASE_				:      int    ,			SCREAMING_SNAKE_CASE_				:      str       ):
        lowerCAmelCase_   :				Optional[int]								= DPTModel(config=SCREAMING_SNAKE_CASE_       )
        model.to(SCREAMING_SNAKE_CASE_       )
        model.eval()
        lowerCAmelCase_   :				Any								= model(SCREAMING_SNAKE_CASE_       )
        self.parent.assertEqual(result.last_hidden_state.shape    ,			(self.batch_size, self.seq_length, self.hidden_size)       )



      def      SCREAMING_SNAKE_CASE__ (      self				:      int    ,			SCREAMING_SNAKE_CASE_				:      int    ,			SCREAMING_SNAKE_CASE_				:      Union[str, Any]    ,			SCREAMING_SNAKE_CASE_				:      Optional[Any]       ):
        lowerCAmelCase_   :				Dict								= self.num_labels
        lowerCAmelCase_   :				Dict								= DPTForDepthEstimation(SCREAMING_SNAKE_CASE_       )
        model.to(SCREAMING_SNAKE_CASE_       )
        model.eval()
        lowerCAmelCase_   :				Tuple								= model(SCREAMING_SNAKE_CASE_       )
        self.parent.assertEqual(result.predicted_depth.shape    ,			(self.batch_size, self.image_size, self.image_size)       )



      def      SCREAMING_SNAKE_CASE__ (      self				:      str    ,			SCREAMING_SNAKE_CASE_				:      List[Any]    ,			SCREAMING_SNAKE_CASE_				:      Any    ,			SCREAMING_SNAKE_CASE_				:      str       ):
        lowerCAmelCase_   :				Union[str, Any]								= self.num_labels
        lowerCAmelCase_   :				str								= DPTForSemanticSegmentation(SCREAMING_SNAKE_CASE_       )
        model.to(SCREAMING_SNAKE_CASE_       )
        model.eval()
        lowerCAmelCase_   :				List[Any]								= model(SCREAMING_SNAKE_CASE_    ,			labels=SCREAMING_SNAKE_CASE_       )
        self.parent.assertEqual(
            result.logits.shape    ,			(self.batch_size, self.num_labels, self.image_size, self.image_size)       )





      def      SCREAMING_SNAKE_CASE__ (      self				:      Optional[int]       ):
        lowerCAmelCase_   :				Optional[Any]								= self.prepare_config_and_inputs()
        lowerCAmelCase_						,lowerCAmelCase_						,lowerCAmelCase_   :				int								= config_and_inputs
        lowerCAmelCase_   :				List[str]								= {'pixel_values': pixel_values}
        return config, inputs_dict







@require_torch
class       UpperCamelCase__   (	lowercase_,					lowercase_,					unittest.TestCase  ):




      """simple docstring"""




      _SCREAMING_SNAKE_CASE  				=     (DPTModel, DPTForDepthEstimation, DPTForSemanticSegmentation) if is_torch_available() else ()
      _SCREAMING_SNAKE_CASE  				=     (
          {
              """depth-estimation""": DPTForDepthEstimation,
              """feature-extraction""": DPTModel,
              """image-segmentation""": DPTForSemanticSegmentation,
          }
          if is_torch_available()
          else {}
      )

      _SCREAMING_SNAKE_CASE  				=     False
      _SCREAMING_SNAKE_CASE  				=     False
      _SCREAMING_SNAKE_CASE  				=     False



      def      SCREAMING_SNAKE_CASE__ (      self				:      Optional[int]       ):
        lowerCAmelCase_   :				List[str]								= DPTModelTester(self       )
        lowerCAmelCase_   :				int								= ConfigTester(self    ,			config_class=SCREAMING_SNAKE_CASE_    ,			has_text_modality=SCREAMING_SNAKE_CASE_    ,			hidden_size=3_7       )



      def      SCREAMING_SNAKE_CASE__ (      self				:      List[Any]       ):
        self.config_tester.run_common_tests()



      @unittest.skip(reason='DPT does not use inputs_embeds'       )
      def      SCREAMING_SNAKE_CASE__ (      self				:      List[str]       ):
        pass



      def      SCREAMING_SNAKE_CASE__ (      self				:      int       ):
        lowerCAmelCase_						,lowerCAmelCase_   :				str								= self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
          lowerCAmelCase_   :				Any								= model_class(SCREAMING_SNAKE_CASE_       )
          self.assertIsInstance(model.get_input_embeddings()    ,			(nn.Module)       )
          lowerCAmelCase_   :				List[str]								= model.get_output_embeddings()
          self.assertTrue(x is None or isinstance(SCREAMING_SNAKE_CASE_    ,			nn.Linear       )       )



      def      SCREAMING_SNAKE_CASE__ (      self				:      int       ):
        lowerCAmelCase_						,lowerCAmelCase_   :				Dict								= self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
          lowerCAmelCase_   :				str								= model_class(SCREAMING_SNAKE_CASE_       )
          lowerCAmelCase_   :				str								= inspect.signature(model.forward       )
          # signature.parameters is an OrderedDict => so arg_names order is deterministic
          lowerCAmelCase_   :				str								= [*signature.parameters.keys()]

          lowerCAmelCase_   :				List[str]								= ['pixel_values']
          self.assertListEqual(arg_names[:1]    ,			SCREAMING_SNAKE_CASE_       )



      def      SCREAMING_SNAKE_CASE__ (      self				:      Optional[Any]       ):
        lowerCAmelCase_   :				Optional[Any]								= self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE_       )



      def      SCREAMING_SNAKE_CASE__ (      self				:      str       ):
        lowerCAmelCase_   :				Dict								= self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_depth_estimation(*SCREAMING_SNAKE_CASE_       )



      def      SCREAMING_SNAKE_CASE__ (      self				:      List[Any]       ):
        lowerCAmelCase_   :				int								= self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_semantic_segmentation(*SCREAMING_SNAKE_CASE_       )



      def      SCREAMING_SNAKE_CASE__ (      self				:      int       ):
        for model_class in self.all_model_classes:
          if model_class.__name__ == "DPTForDepthEstimation":
            continue

          lowerCAmelCase_						,lowerCAmelCase_   :				Dict								= self.model_tester.prepare_config_and_inputs_for_common()
          lowerCAmelCase_   :				Tuple								= True

          if model_class in get_values(SCREAMING_SNAKE_CASE_       ):
            continue

          lowerCAmelCase_   :				int								= model_class(SCREAMING_SNAKE_CASE_       )
          model.to(SCREAMING_SNAKE_CASE_       )
          model.train()
          lowerCAmelCase_   :				Dict								= self._prepare_for_class(SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_    ,			return_labels=SCREAMING_SNAKE_CASE_       )
          lowerCAmelCase_   :				str								= model(**SCREAMING_SNAKE_CASE_       ).loss
          loss.backward()



      def      SCREAMING_SNAKE_CASE__ (      self				:      int       ):
        for model_class in self.all_model_classes:
          if model_class.__name__ == "DPTForDepthEstimation":
            continue

          lowerCAmelCase_						,lowerCAmelCase_   :				List[str]								= self.model_tester.prepare_config_and_inputs_for_common()
          lowerCAmelCase_   :				int								= False
          lowerCAmelCase_   :				Union[str, Any]								= True

          if model_class in get_values(SCREAMING_SNAKE_CASE_       ) or not model_class.supports_gradient_checkpointing:
            continue
          lowerCAmelCase_   :				Dict								= model_class(SCREAMING_SNAKE_CASE_       )
          model.to(SCREAMING_SNAKE_CASE_       )
          model.gradient_checkpointing_enable()
          model.train()
          lowerCAmelCase_   :				List[str]								= self._prepare_for_class(SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_    ,			return_labels=SCREAMING_SNAKE_CASE_       )
          lowerCAmelCase_   :				Union[str, Any]								= model(**SCREAMING_SNAKE_CASE_       ).loss
          loss.backward()



      def      SCREAMING_SNAKE_CASE__ (      self				:      Tuple       ):
        lowerCAmelCase_						,lowerCAmelCase_   :				Dict								= self.model_tester.prepare_config_and_inputs_for_common()

        lowerCAmelCase_   :				List[Any]								= _config_zero_init(SCREAMING_SNAKE_CASE_       )
        for model_class in self.all_model_classes:
          lowerCAmelCase_   :				Optional[int]								= model_class(config=SCREAMING_SNAKE_CASE_       )
          # Skip the check for the backbone
          lowerCAmelCase_   :				Union[str, Any]								= []
          for name, module in model.named_modules():
            if module.__class__.__name__ == "DPTViTHybridEmbeddings":
              lowerCAmelCase_   :				List[Any]								= [F"{name}.{key}" for key in module.state_dict().keys()]
              break

          for name, param in model.named_parameters():
            if param.requires_grad:
              if name in backbone_params:
                continue
              self.assertIn(
                  ((param.data.mean() * 1E9).round() / 1E9).item()    ,			[0.0, 1.0]    ,			msg=F"Parameter {name} of model {model_class} seems not properly initialized"    ,			)



      @unittest.skip('Will be fixed soon by reducing the size of the model used for common tests.'       )
      def      SCREAMING_SNAKE_CASE__ (      self				:      Dict       ):
        pass



      @slow
      def      SCREAMING_SNAKE_CASE__ (      self				:      List[str]       ):
        for model_name in DPT_PRETRAINED_MODEL_ARCHIVE_LIST[1:]:
          lowerCAmelCase_   :				Tuple								= DPTModel.from_pretrained(SCREAMING_SNAKE_CASE_       )
          self.assertIsNotNone(SCREAMING_SNAKE_CASE_       )





      def      SCREAMING_SNAKE_CASE__ (      self				:      Optional[int]       ):
        # We do this test only for DPTForDepthEstimation since it is the only model that uses readout_type
        lowerCAmelCase_						,lowerCAmelCase_   :				List[str]								= self.model_tester.prepare_config_and_inputs_for_common()
        lowerCAmelCase_   :				Dict								= 'add'
        with self.assertRaises(SCREAMING_SNAKE_CASE_       ):
          lowerCAmelCase_   :				Optional[int]								= DPTForDepthEstimation(SCREAMING_SNAKE_CASE_       )




def 				UpperCamelCase_    (    )						->   Tuple:

  """simple docstring"""




  lowerCAmelCase_   :				List[str]								= Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'						)
  return image







@require_torch
@require_vision
@slow
class       UpperCamelCase__   (	unittest.TestCase  ):




      """simple docstring"""




      def      SCREAMING_SNAKE_CASE__ (      self				:      List[str]       ):
        lowerCAmelCase_   :				str								= DPTImageProcessor.from_pretrained('Intel/dpt-hybrid-midas'       )
        lowerCAmelCase_   :				str								= DPTForDepthEstimation.from_pretrained('Intel/dpt-hybrid-midas'       ).to(SCREAMING_SNAKE_CASE_       )

        lowerCAmelCase_   :				int								= prepare_img()
        lowerCAmelCase_   :				Optional[int]								= image_processor(images=SCREAMING_SNAKE_CASE_    ,			return_tensors='pt'       ).to(SCREAMING_SNAKE_CASE_       )

        # forward pass
        with torch.no_grad():
          lowerCAmelCase_   :				int								= model(**SCREAMING_SNAKE_CASE_       )
          lowerCAmelCase_   :				int								= outputs.predicted_depth

        # verify the predicted depth
        lowerCAmelCase_   :				List[Any]								= torch.Size((1, 3_8_4, 3_8_4)       )
        self.assertEqual(predicted_depth.shape    ,			SCREAMING_SNAKE_CASE_       )

        lowerCAmelCase_   :				Dict								= torch.tensor(
            [[[5.64_37, 5.61_46, 5.65_11], [5.43_71, 5.56_49, 5.59_58], [5.52_15, 5.51_84, 5.52_93]]]       ).to(SCREAMING_SNAKE_CASE_       )

        self.assertTrue(torch.allclose(outputs.predicted_depth[:3, :3, :3] / 1_0_0    ,			SCREAMING_SNAKE_CASE_    ,			atol=1E-4       )       )










"""simple docstring"""







import argparse
import logging
import os

import datasets
import tensorflow as tf

from transformers import AutoTokenizer


lowercase__		:				List[Any]          =	logging.getLogger(__name__)

def 				UpperCamelCase_    (    )						->   Dict:

  """simple docstring"""




  lowerCAmelCase_   :				Tuple								= 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_   :				List[Any]								= parser.parse_args()
  return args

def 				UpperCamelCase_    (    lowerCAmelCase__      :				int						)						->   Union[str, Any]:

  """simple docstring"""




  def fn(lowerCAmelCase__      :				Optional[Any]						):
    return tokenizer(examples['text']						)

  return fn

def 				UpperCamelCase_    (    lowerCAmelCase__      :				Tuple						)						->   Dict:

  """simple docstring"""




  lowerCAmelCase_   :				int								= []
  for i in range(len(tokenized_data['input_ids']						)						):
    lowerCAmelCase_   :				Tuple								= {
        '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_   :				Union[str, Any]								= tf.train.Features(feature=lowerCAmelCase__						)
    lowerCAmelCase_   :				Dict								= tf.train.Example(features=lowerCAmelCase__						)
    lowerCAmelCase_   :				Optional[int]								= example.SerializeToString()
    records.append(lowerCAmelCase__						)
  return records

def 				UpperCamelCase_    (    lowerCAmelCase__      :				Optional[Any]						)						->   Optional[int]:

  """simple docstring"""




  lowerCAmelCase_   :				str								= datasets.load_dataset(args.dataset_name    ,   args.dataset_config    ,   split=args.split						)

  if args.limit is not None:
    lowerCAmelCase_   :				Tuple								= min(len(lowerCAmelCase__						)    ,   args.limit						)
    lowerCAmelCase_   :				Any								= dataset.select(range(lowerCAmelCase__						)						)
    print(f"Limiting the dataset to {args.limit} entries."						)

  lowerCAmelCase_   :				List[str]								= 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_   :				int								= os.path.join(args.output_dir    ,   args.split						)
    if not os.path.exists(lowerCAmelCase__						):
      os.makedirs(lowerCAmelCase__						)
  else:
    lowerCAmelCase_   :				Dict								= os.path.join(args.output_dir    ,   args.split						)

  # Tokenize the whole dataset at once.
  lowerCAmelCase_   :				Dict								= tokenize_function(lowerCAmelCase__						)
  lowerCAmelCase_   :				List[Any]								= 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[Any]						):
    # Concatenate all texts.
    lowerCAmelCase_   :				int								= {k: sum(examples[k]    ,   []						) for k in examples.keys()}
    lowerCAmelCase_   :				Any								= 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_   :				Union[str, Any]								= (total_length // args.max_length) * args.max_length
    # Split by chunks of max_len.
    lowerCAmelCase_   :				Optional[Any]								= {
        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_   :				Optional[int]								= dataset_tokenized.map(lowerCAmelCase__    ,   batched=lowerCAmelCase__    ,   batch_size=1000    ,   num_proc=4						)

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

    with tf.io.TFRecordWriter(lowerCAmelCase__						) as out_file:
      for i in range(len(lowerCAmelCase__						)						):
        lowerCAmelCase_   :				Dict								= 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__":
     lowercase__		:				int          =	parse_args()
     main(args)





