code: def main(args):
      if not args.no_tpu:
          tpu = initialize_tpu(args)
          strategy = tf.distribute.TPUStrategy(tpu)
      else:
          strategy = tf.distribute.OneDeviceStrategy(device="/gpu:0")
  
      if args.bfloat16:
          tf.keras.mixed_precision.set_global_policy("mixed_bfloat16")
  
      tokenizer = AutoTokenizer.from_pretrained(args.tokenizer)
      config = AutoConfig.from_pretrained(args.pretrained_model_config)
      config.vocab_size = tokenizer.vocab_size
  
      training_records = tf.io.gfile.glob(
          os.path.join(args.train_dataset, "*.tfrecord")
      )
      if not training_records:
          raise ValueError(
              f"No .tfrecord files found in {args.train_dataset}."
          )
      eval_records = tf.io.gfile.glob(
          os.path.join(args.eval_dataset, "*.tfrecord")
      )
      if not eval_records:
          raise ValueError(
              f"No .tfrecord files found in {args.eval_dataset}."
          )
  
      num_train_samples = count_samples(training_records)
  
      steps_per_epoch = num_train_samples // (
          args.per_replica_batch_size * strategy.num_replicas_in_sync
      )
      total_train_steps = steps_per_epoch * args.num_epochs
  
      with strategy.scope():
          model = TFAutoModelForMaskedLM.from_config(config)
          model(
              model.dummy_inputs
          )  # Pass some dummy inputs through the model to ensure all the weights are built
          optimizer, schedule = create_optimizer(
              num_train_steps=total_train_steps,
              num_warmup_steps=total_train_steps // 20,
              init_lr=args.learning_rate,
              weight_decay_rate=args.weight_decay_rate,
          )
  
          # Transformers models compute the right loss for their task by default when labels are passed, and will
          # use this for training unless you specify your own loss function in compile().
          model.compile(optimizer=optimizer, metrics=["accuracy"])
  
      def decode_fn(example):
          features = {
              "input_ids": tf.io.FixedLenFeature(
                  dtype=tf.int64, shape=(args.max_length,)
              ),
              "attention_mask": tf.io.FixedLenFeature(
                  dtype=tf.int64, shape=(args.max_length,)
              ),
          }
          return tf.io.parse_single_example(example, features)
  
      # Many of the data collators in Transformers are TF-compilable when return_tensors == "tf", so we can
      # use their methods in our data pipeline.
      data_collator = DataCollatorForLanguageModeling(
          tokenizer=tokenizer,
          mlm_probability=args.mlm_probability,
          mlm=True,
          return_tensors="tf",
      )
  
      def mask_with_collator(batch):
          # TF really needs an isin() function
          special_tokens_mask = (
              ~tf.cast(batch["attention_mask"], tf.bool)
              | (batch["input_ids"] == tokenizer.cls_token_id)
              | (batch["input_ids"] == tokenizer.sep_token_id)
          )
          (
              batch["input_ids"],
              batch["labels"],
          ) = data_collator.tf_mask_tokens(
              batch["input_ids"],
              vocab_size=len(tokenizer),
              mask_token_id=tokenizer.mask_token_id,
              special_tokens_mask=special_tokens_mask,
          )
          return batch
  
      batch_size = (
          args.per_replica_batch_size * strategy.num_replicas_in_sync
      )
  
      train_dataset = prepare_dataset(
          training_records,
          decode_fn=decode_fn,
          mask_fn=mask_with_collator,
          batch_size=batch_size,
          shuffle=True,
          shuffle_buffer_size=args.shuffle_buffer_size,
      )
  
      eval_dataset = prepare_dataset(
          eval_records,
          decode_fn=decode_fn,
          mask_fn=mask_with_collator,
          batch_size=batch_size,
          shuffle=False,
      )
  
      callbacks = []
      if args.hub_model_id:
          callbacks.append(
              PushToHubCallback(
                  output_dir=args.output_dir,
                  hub_model_id=args.hub_model_id,
                  tokenizer=tokenizer,
              )
          )
  
      model.fit(
          train_dataset,
          validation_data=eval_dataset,
          epochs=args.num_epochs,
          callbacks=callbacks,
      )
  
      model.save_pretrained(args.output_dir)
  
file_path: transformers/examples/tensorflow/language-modeling-tpu/run_mlm.py
imports: import argparse
  import logging
  import os
  import re
  import tensorflow
  from transformers import AutoConfig, AutoTokenizer, DataCollatorForLanguageModeling, PushToHubCallback, TFAutoModelForMaskedLM, create_optimizer
label_desc: Use image node_id 5 for calling a global function with example usage: main(args) and returns: tf, batch