code: def main(args):
      dataset = datasets.load_dataset(
          args.dataset_name, args.dataset_config, split=args.split
      )
  
      if args.limit is not None:
          max_samples = min(len(dataset), args.limit)
          dataset = dataset.select(range(max_samples))
          print(f"Limiting the dataset to {args.limit} entries.")
  
      tokenizer = 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)
          split_dir = os.path.join(args.output_dir, args.split)
          if not os.path.exists(split_dir):
              os.makedirs(split_dir)
      else:
          split_dir = os.path.join(args.output_dir, args.split)
  
      # Tokenize the whole dataset at once.
      tokenize_fn = tokenize_function(tokenizer)
      dataset_tokenized = dataset.map(
          tokenize_fn, batched=True, 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(examples):
          # Concatenate all texts.
          concatenated_examples = {
              k: sum(examples[k], []) for k in examples.keys()
          }
          total_length = 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 🫀
          total_length = (
              total_length // args.max_length
          ) * args.max_length
          # Split by chunks of max_len.
          result = {
              k: [
                  t[i : i + args.max_length]
                  for i in range(0, total_length, args.max_length)
              ]
              for k, t in concatenated_examples.items()
          }
          return result
  
      grouped_dataset = dataset_tokenized.map(
          group_texts, batched=True, batch_size=1000, num_proc=4
      )
  
      shard_count = 0
      total_records = 0
      for shard in range(0, len(grouped_dataset), args.shard_size):
          dataset_snapshot = grouped_dataset[
              shard : shard + args.shard_size
          ]
          records_containing = len(dataset_snapshot["input_ids"])
          filename = os.path.join(
              split_dir,
              f"dataset-{shard_count}-{records_containing}.tfrecord",
          )
          serialized_examples = get_serialized_examples(
              dataset_snapshot
          )
  
          with tf.io.TFRecordWriter(filename) as out_file:
              for i in range(len(serialized_examples)):
                  example = serialized_examples[i]
                  out_file.write(example)
              print(
                  "Wrote file {} containing {} records".format(
                      filename, records_containing
                  )
              )
  
          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=f)
  
file_path: transformers/examples/tensorflow/language-modeling-tpu/prepare_tfrecord_shards.py
imports: import argparse
  import logging
  import os
  import datasets
  import tensorflow
  from transformers import AutoTokenizer
label_desc: Use image node_id 4 for calling a global function with example usage: main(args) and returns: result