# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Utilities used for data preparation."""

import collections
import json
import os

from absl import logging

import numpy as np
import tensorflow as tf, tf_keras

special_symbols = {
    "<unk>": 0,
    "<s>": 1,
    "</s>": 2,
    "<cls>": 3,
    "<sep>": 4,
    "<pad>": 5,
    "<mask>": 6,
    "<eod>": 7,
    "<eop>": 8,
}

VOCAB_SIZE = 32000
UNK_ID = special_symbols["<unk>"]
CLS_ID = special_symbols["<cls>"]
SEP_ID = special_symbols["<sep>"]
MASK_ID = special_symbols["<mask>"]
EOD_ID = special_symbols["<eod>"]
SEG_ID_P = 0
SEG_ID_Q = 1
SEG_ID_CLS = 2
SEG_ID_PAD = 3

OnlineMaskingConfig = collections.namedtuple("OnlineMaskingConfig", [
    "sample_strategy", "max_num_tokens", "min_num_tokens", "max_num_words",
    "min_num_words"
])


def file_based_input_fn_builder(input_file, name_to_features, batch_size,
                                is_training):
  """Creates an `input_fn` closure."""

  logging.info("Input tfrecord file %s", input_file)

  def _decode_record(record, name_to_features):
    """Decodes a record to a TensorFlow example."""
    example = tf.io.parse_single_example(record, name_to_features)

    # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
    # So cast all int64 to int32.
    for name in list(example.keys()):
      t = example[name]
      if t.dtype == tf.int64:
        t = tf.cast(t, tf.int32)
      example[name] = t

    return example

  def input_fn():
    """Returns dataset for training/evaluation."""
    num_threads = 8
    if isinstance(input_file, str):
      d = tf.data.TFRecordDataset(input_file)
      # For training, we want a lot of parallel reading and shuffling.
      # For eval, we want no shuffling and parallel reading doesn't matter.
      if is_training:
        d = d.shuffle(2048)
        d = d.repeat()
    else:
      cycle_length = min(num_threads, len(input_file))
      d = tf.data.Dataset.from_tensor_slices(input_file)
      # file level shuffle
      d = d.shuffle(len(input_file)).repeat()

      d = d.interleave(
          tf.data.TFRecordDataset,
          cycle_length=cycle_length)

      if is_training:
        # sample level shuffle
        d = d.shuffle(buffer_size=2048)
    d = d.map(
        lambda record: _decode_record(record, name_to_features),
        num_parallel_calls=tf.data.experimental.AUTOTUNE)
    d = d.batch(batch_size, drop_remainder=is_training)

    # When `input_file` is a path to a single file or a list
    # containing a single path, disable auto sharding so that
    # same input file is sent to all workers.
    if isinstance(input_file, str) or len(input_file) == 1:
      options = tf.data.Options()
      options.experimental_distribute.auto_shard_policy = (
          tf.data.experimental.AutoShardPolicy.OFF)
      d = d.with_options(options)

    d = d.prefetch(tf.data.experimental.AUTOTUNE)
    return d

  return input_fn


def create_classification_dataset(file_path, seq_length, batch_size,
                                  is_training):
  """Creates input dataset from (tf)records files for pretraining."""
  name_to_features = {
      "input_ids": tf.io.FixedLenFeature([seq_length], tf.int64),
      "input_mask": tf.io.FixedLenFeature([seq_length], tf.float32),
      "segment_ids": tf.io.FixedLenFeature([seq_length], tf.int64),
      "label_ids": tf.io.FixedLenFeature([], tf.int64),
      "is_real_example": tf.io.FixedLenFeature([], tf.int64),
  }

  input_fn = file_based_input_fn_builder(file_path, name_to_features,
                                         batch_size, is_training)
  dataset = input_fn()
  return dataset


def create_squad_dataset(file_path, seq_length, batch_size, is_training):
  """Creates input dataset from (tf)records files for pretraining."""
  name_to_features = {
      "unique_ids": tf.io.FixedLenFeature([], tf.int64),
      "input_ids": tf.io.FixedLenFeature([seq_length], tf.int64),
      "input_mask": tf.io.FixedLenFeature([seq_length], tf.float32),
      "segment_ids": tf.io.FixedLenFeature([seq_length], tf.int64),
      "cls_index": tf.io.FixedLenFeature([], tf.int64),
      "p_mask": tf.io.FixedLenFeature([seq_length], tf.float32)
  }

  if is_training:
    name_to_features["start_positions"] = tf.io.FixedLenFeature([], tf.int64)
    name_to_features["end_positions"] = tf.io.FixedLenFeature([], tf.int64)
    name_to_features["is_impossible"] = tf.io.FixedLenFeature([], tf.float32)

  input_fn = file_based_input_fn_builder(file_path, name_to_features,
                                         batch_size, is_training)
  dataset = input_fn()
  return dataset


def get_input_iterator(input_fn, strategy):
  """Returns distributed dataset iterator."""

  # When training with TPU pods, datasets needs to be cloned across
  # workers. Since Dataset instance cannot be cloned in eager mode, we instead
  # pass callable that returns a dataset.
  input_data = input_fn()
  if callable(input_data):
    iterator = iter(strategy.distribute_datasets_from_function(input_data))
  else:
    iterator = iter(strategy.experimental_distribute_dataset(input_data))
  return iterator


def get_classification_input_data(batch_size, seq_len, strategy, is_training,
                                  file_path):
  """Returns input dataset from input file string."""

  # When using TPU pods, we need to clone dataset across
  # workers and need to pass in function that returns the dataset rather
  # than passing dataset instance itself.
  use_dataset_fn = isinstance(strategy, tf.distribute.TPUStrategy)
  if use_dataset_fn:
    if batch_size % strategy.num_replicas_in_sync != 0:
      raise ValueError(
          "Batch size must be divisible by number of replicas : {}".format(
              strategy.num_replicas_in_sync))

    # As auto rebatching is not supported in
    # `distribute_datasets_from_function()` API, which is
    # required when cloning dataset to multiple workers in eager mode,
    # we use per-replica batch size.
    batch_size = int(batch_size / strategy.num_replicas_in_sync)

  def _dataset_fn(ctx=None):
    del ctx

    train_dataset = create_classification_dataset(
        file_path=file_path,
        seq_length=seq_len,
        batch_size=batch_size,
        is_training=is_training)
    return train_dataset

  return _dataset_fn if use_dataset_fn else _dataset_fn()


def get_squad_input_data(batch_size, seq_len, q_len, strategy, is_training,
                         file_path):
  """Returns input dataset from input file string."""

  # When using TPU pods, we need to clone dataset across
  # workers and need to pass in function that returns the dataset rather
  # than passing dataset instance itself.
  use_dataset_fn = isinstance(strategy, tf.distribute.TPUStrategy)
  if use_dataset_fn:
    if batch_size % strategy.num_replicas_in_sync != 0:
      raise ValueError(
          "Batch size must be divisible by number of replicas : {}".format(
              strategy.num_replicas_in_sync))

    # As auto rebatching is not supported in
    # `distribute_datasets_from_function()` API, which is
    # required when cloning dataset to multiple workers in eager mode,
    # we use per-replica batch size.
    batch_size = int(batch_size / strategy.num_replicas_in_sync)

  if is_training:
    input_glob = os.path.join(
        file_path,
        "spiece.model.*.slen-{}.qlen-{}.train.tf_record".format(seq_len, q_len))

    global_input_paths = tf.io.gfile.glob(input_glob)
  else:
    global_input_paths = file_path

  def _dataset_fn(ctx=None):
    del ctx

    train_dataset = create_squad_dataset(
        file_path=global_input_paths,
        seq_length=seq_len,
        batch_size=batch_size,
        is_training=is_training)
    return train_dataset

  return _dataset_fn if use_dataset_fn else _dataset_fn()


def _idx_pair_to_mask(beg_indices, end_indices, inputs, tgt_len, num_predict):
  """Turn beg and end indices into actual mask."""
  non_func_mask = tf.logical_and(
      tf.not_equal(inputs, SEP_ID), tf.not_equal(inputs, CLS_ID))
  all_indices = tf.where(non_func_mask, tf.range(tgt_len, dtype=tf.int64),
                         tf.constant(-1, shape=[tgt_len], dtype=tf.int64))
  candidate_matrix = tf.cast(
      tf.logical_and(all_indices[None, :] >= beg_indices[:, None],
                     all_indices[None, :] < end_indices[:, None]), tf.float32)
  cumsum_matrix = tf.reshape(
      tf.cumsum(tf.reshape(candidate_matrix, [-1])), [-1, tgt_len])
  masked_matrix = tf.cast(cumsum_matrix <= num_predict, tf.float32)
  target_mask = tf.reduce_sum(candidate_matrix * masked_matrix, axis=0)
  is_masked = tf.cast(target_mask, tf.bool)

  return is_masked, target_mask


def _word_span_mask(inputs, tgt_len, num_predict, min_num_words, max_num_words,
                    boundary):
  """Sample whole word spans as prediction targets."""
  # Note: 1.2 is the token-to-word ratio
  mask_alpha = tgt_len / num_predict / 1.2
  round_to_int = lambda x: tf.cast(tf.round(x), tf.int64)

  # Sample span lengths from a zipf distribution
  span_len_seq = np.arange(min_num_words, max_num_words + 1)
  probs = np.array([1.0 / (i + 1) for i in span_len_seq])
  probs /= np.sum(probs)
  logits = tf.constant(np.log(probs), dtype=tf.float32)

  # Sample `num_predict` words here: note that this is over sampling
  span_lens = tf.random.categorical(
      logits=logits[None],
      num_samples=num_predict,
      dtype=tf.int64,
  )[0] + min_num_words

  # Sample the ratio [0.0, 1.0) of left context lengths
  span_lens_float = tf.cast(span_lens, tf.float32)
  left_ratio = tf.random.uniform(shape=[num_predict], minval=0.0, maxval=1.0)
  left_ctx_len = left_ratio * span_lens_float * (mask_alpha - 1)

  left_ctx_len = round_to_int(left_ctx_len)
  right_offset = round_to_int(span_lens_float * mask_alpha) - left_ctx_len

  beg_indices = (
      tf.cumsum(left_ctx_len) + tf.cumsum(right_offset, exclusive=True))
  end_indices = beg_indices + span_lens

  # Remove out of range indices
  max_boundary_index = tf.cast(tf.shape(boundary)[0] - 1, tf.int64)
  valid_idx_mask = end_indices < max_boundary_index
  beg_indices = tf.boolean_mask(beg_indices, valid_idx_mask)
  end_indices = tf.boolean_mask(end_indices, valid_idx_mask)

  beg_indices = tf.gather(boundary, beg_indices)
  end_indices = tf.gather(boundary, end_indices)

  # Shuffle valid indices
  num_valid = tf.cast(tf.shape(beg_indices)[0], tf.int64)
  order = tf.random.shuffle(tf.range(num_valid, dtype=tf.int64))
  beg_indices = tf.gather(beg_indices, order)
  end_indices = tf.gather(end_indices, order)

  return _idx_pair_to_mask(beg_indices, end_indices, inputs, tgt_len,
                           num_predict)


def _token_span_mask(inputs, tgt_len, num_predict, min_num_tokens,
                     max_num_tokens):
  """Sample token spans as prediction targets."""
  mask_alpha = tgt_len / num_predict
  round_to_int = lambda x: tf.cast(tf.round(x), tf.int64)

  # Sample span lengths from a zipf distribution
  span_len_seq = np.arange(min_num_tokens, max_num_tokens + 1)
  probs = np.array([1.0 / (i + 1) for i in span_len_seq])

  probs /= np.sum(probs)
  logits = tf.constant(np.log(probs), dtype=tf.float32)
  span_lens = tf.random.categorical(
      logits=logits[None],
      num_samples=num_predict,
      dtype=tf.int64,
  )[0] + min_num_tokens

  # Sample the ratio [0.0, 1.0) of left context lengths
  span_lens_float = tf.cast(span_lens, tf.float32)
  left_ratio = tf.random.uniform(shape=[num_predict], minval=0.0, maxval=1.0)
  left_ctx_len = left_ratio * span_lens_float * (mask_alpha - 1)
  left_ctx_len = round_to_int(left_ctx_len)

  # Compute the offset from left start to the right end
  right_offset = round_to_int(span_lens_float * mask_alpha) - left_ctx_len

  # Get the actual begin and end indices
  beg_indices = (
      tf.cumsum(left_ctx_len) + tf.cumsum(right_offset, exclusive=True))
  end_indices = beg_indices + span_lens

  # Remove out of range indices
  valid_idx_mask = end_indices < tgt_len
  beg_indices = tf.boolean_mask(beg_indices, valid_idx_mask)
  end_indices = tf.boolean_mask(end_indices, valid_idx_mask)

  # Shuffle valid indices
  num_valid = tf.cast(tf.shape(beg_indices)[0], tf.int64)
  order = tf.random.shuffle(tf.range(num_valid, dtype=tf.int64))
  beg_indices = tf.gather(beg_indices, order)
  end_indices = tf.gather(end_indices, order)

  return _idx_pair_to_mask(beg_indices, end_indices, inputs, tgt_len,
                           num_predict)


def _whole_word_mask(inputs, tgt_len, num_predict, boundary):
  """Sample whole words as prediction targets."""
  pair_indices = tf.concat([boundary[:-1, None], boundary[1:, None]], axis=1)
  cand_pair_indices = tf.random.shuffle(pair_indices)[:num_predict]
  beg_indices = cand_pair_indices[:, 0]
  end_indices = cand_pair_indices[:, 1]

  return _idx_pair_to_mask(beg_indices, end_indices, inputs, tgt_len,
                           num_predict)


def _single_token_mask(inputs, tgt_len, num_predict):
  """Sample individual tokens as prediction targets."""
  all_indices = tf.range(tgt_len, dtype=tf.int64)
  non_func_mask = tf.logical_and(
      tf.not_equal(inputs, SEP_ID), tf.not_equal(inputs, CLS_ID))
  non_func_indices = tf.boolean_mask(all_indices, non_func_mask)

  masked_pos = tf.random.shuffle(non_func_indices)
  masked_pos = tf.sort(masked_pos[:num_predict])
  target_mask = tf.sparse_to_dense(
      sparse_indices=masked_pos,
      output_shape=[tgt_len],
      sparse_values=1.0,
      default_value=0.0)

  is_masked = tf.cast(target_mask, tf.bool)

  return is_masked, target_mask


def _online_sample_masks(inputs,
                         tgt_len,
                         num_predict,
                         online_masking_config,
                         boundary=None):
  """Sample target positions to predict."""
  logging.info("Online sample with strategy: `%s`.",
               online_masking_config.sample_strategy)
  if online_masking_config.sample_strategy == "single_token":
    return _single_token_mask(inputs, tgt_len, num_predict)
  elif online_masking_config.sample_strategy == "whole_word":
    assert boundary is not None, "whole word sampling requires `boundary`"
    return _whole_word_mask(inputs, tgt_len, num_predict, boundary)
  elif online_masking_config.sample_strategy == "token_span":
    return _token_span_mask(inputs, tgt_len, num_predict,
                            online_masking_config.min_num_tokens,
                            online_masking_config.max_num_tokens)
  elif online_masking_config.sample_strategy == "word_span":
    assert boundary is not None, "word span sampling requires `boundary`"
    return _word_span_mask(inputs, tgt_len, num_predict,
                           online_masking_config.min_num_words,
                           online_masking_config.max_num_words, boundary)
  else:
    raise NotImplementedError


def create_pretrain_dataset(file_names,
                            bsz_per_core,
                            seq_len,
                            reuse_len,
                            perm_size,
                            leak_ratio,
                            online_masking_config,
                            num_predict=None,
                            input_pipeline_context=None):
  """Creates pretrain dataset."""

  def parser(record):
    """Function used to parse tfrecord."""

    record_spec = {
        "input": tf.io.FixedLenFeature([seq_len], tf.int64),
        "seg_id": tf.io.FixedLenFeature([seq_len], tf.int64),
        "label": tf.io.FixedLenFeature([1], tf.int64),
    }

    if online_masking_config.sample_strategy in ["whole_word", "word_span"]:
      logging.info("Add `boundary` spec for %s",
                   online_masking_config.sample_strategy)
      record_spec["boundary"] = tf.io.VarLenFeature(tf.int64)

    # retrieve serialized example
    example = tf.io.parse_single_example(
        serialized=record, features=record_spec)

    inputs = example.pop("input")
    if online_masking_config.sample_strategy in ["whole_word", "word_span"]:
      boundary = tf.sparse.to_dense(example.pop("boundary"))
    else:
      boundary = None
    is_masked, _ = _online_sample_masks(
        inputs, seq_len, num_predict, online_masking_config, boundary=boundary)

    if reuse_len > 0:
      ##### Use memory
      # permutate the reuse and non-reuse parts separately
      non_reuse_len = seq_len - reuse_len
      assert reuse_len % perm_size == 0 and non_reuse_len % perm_size == 0

      # Creates permutation mask and target mask for the first reuse_len tokens.
      # The tokens in this part are reused from the last sequence.
      perm_mask_0, target_mask_0, input_k_0, input_q_0 = _local_perm(
          inputs[:reuse_len], is_masked[:reuse_len], perm_size, reuse_len,
          leak_ratio)

      # Creates permutation mask and target mask for the rest of tokens in
      # current example, which are concatentation of two new segments.
      perm_mask_1, target_mask_1, input_k_1, input_q_1 = _local_perm(
          inputs[reuse_len:], is_masked[reuse_len:], perm_size, non_reuse_len,
          leak_ratio)

      perm_mask_0 = tf.concat(
          [perm_mask_0, tf.ones([reuse_len, non_reuse_len])], axis=1)
      perm_mask_1 = tf.concat(
          [tf.zeros([non_reuse_len, reuse_len]), perm_mask_1], axis=1)
      perm_mask = tf.concat([perm_mask_0, perm_mask_1], axis=0)
      target_mask = tf.concat([target_mask_0, target_mask_1], axis=0)
      input_k = tf.concat([input_k_0, input_k_1], axis=0)
      input_q = tf.concat([input_q_0, input_q_1], axis=0)
    else:
      ##### Do not use memory
      assert seq_len % perm_size == 0
      # permutate the entire sequence together
      perm_mask, target_mask, input_k, input_q = _local_perm(
          inputs, is_masked, perm_size, seq_len, leak_ratio)

    # reshape back to fixed shape
    example["perm_mask"] = tf.reshape(perm_mask, [seq_len, seq_len])
    example["input_ids"] = tf.reshape(input_k, [seq_len])
    example["input_q"] = tf.reshape(input_q, [seq_len])

    # Directly use raw inputs as the target
    target = inputs

    if num_predict is not None:
      indices = tf.range(seq_len, dtype=tf.int64)
      bool_target_mask = tf.cast(target_mask, tf.bool)
      indices = tf.boolean_mask(indices, bool_target_mask)

      ##### extra padding due to CLS/SEP introduced after prepro
      actual_num_predict = tf.shape(indices)[0]
      pad_len = num_predict - actual_num_predict

      ##### target_mapping
      target_mapping = tf.one_hot(indices, seq_len, dtype=tf.float32)
      paddings = tf.zeros([pad_len, seq_len], dtype=target_mapping.dtype)
      target_mapping = tf.concat([target_mapping, paddings], axis=0)
      example["target_mapping"] = tf.reshape(target_mapping,
                                             [num_predict, seq_len])

      ##### target
      target = tf.boolean_mask(target, bool_target_mask)
      paddings = tf.zeros([pad_len], dtype=target.dtype)
      target = tf.concat([target, paddings], axis=0)
      example["target"] = tf.reshape(target, [num_predict])

      ##### target mask
      target_mask = tf.concat([
          tf.ones([actual_num_predict], dtype=tf.float32),
          tf.zeros([pad_len], dtype=tf.float32)
      ],
                              axis=0)
      example["target_mask"] = tf.reshape(target_mask, [num_predict])
    else:
      example["target"] = tf.reshape(target, [seq_len])
      example["target_mask"] = tf.reshape(target_mask, [seq_len])

    for key in list(example.keys()):
      val = example[key]
      if tf_keras.backend.is_sparse(val):
        val = tf.sparse.to_dense(val)
      if val.dtype == tf.int64:
        val = tf.cast(val, tf.int32)

      example[key] = val

    for k, v in example.items():
      logging.info("%s: %s", k, v)

    return example

  dataset = parse_files_to_dataset(
      parser=parser,
      file_paths=file_names,
      bsz_per_core=bsz_per_core,
      sequential=reuse_len > 0,
      input_pipeline_context=input_pipeline_context)

  return dataset


def format_filename(prefix,
                    suffix,
                    bsz_per_host,
                    seq_len,
                    reuse_len=None,
                    uncased=False):
  """Generates input file name pattern."""
  if reuse_len is not None and reuse_len > 0:
    reuse_str = "reuse-{}.".format(reuse_len)
    bsz_str = "hostbsz-{}.".format(bsz_per_host)
  else:
    reuse_str = ""
    bsz_str = ""

  if not uncased:
    case_str = ""
  else:
    case_str = "uncased."

  file_name = "{}.seq-{}.{}{}{}{}".format(prefix, seq_len, reuse_str, bsz_str,
                                          case_str, suffix)

  return file_name


def get_pretrain_input_data(batch_size,
                            seq_len,
                            strategy,
                            file_path,
                            reuse_len,
                            perm_size,
                            leak_ratio,
                            num_predict,
                            uncased,
                            online_masking_config,
                            num_hosts=1):
  """Returns input dataset from input file string."""

  # When using TPU pods, we need to clone dataset across
  # workers and need to pass in function that returns the dataset rather
  # than passing dataset instance itself.
  use_dataset_fn = isinstance(strategy, tf.distribute.TPUStrategy)
  split = "train"
  bsz_per_host = int(batch_size / num_hosts)
  record_glob_base = format_filename(
      prefix="meta.{}.pass-*".format(split),
      suffix="json*",
      bsz_per_host=bsz_per_host,
      seq_len=seq_len,
      reuse_len=reuse_len,
      uncased=uncased)

  def _get_num_batch(info):
    if "num_batch" in info:
      return info["num_batch"]
    elif "num_example" in info:
      return info["num_example"] / bsz_per_host
    else:
      raise ValueError("Do not have sample info.")

  if use_dataset_fn:
    if batch_size % strategy.num_replicas_in_sync != 0:
      raise ValueError(
          "Batch size must be divisible by number of replicas : {}".format(
              strategy.num_replicas_in_sync))

    # As auto rebatching is not supported in
    # `distribute_datasets_from_function()` API, which is
    # required when cloning dataset to multiple workers in eager mode,
    # we use per-replica batch size.
    batch_size = int(batch_size / strategy.num_replicas_in_sync)

  record_info = {"num_batch": 0, "filenames": []}

  tfrecord_dirs = file_path.split(",")
  logging.info("Use the following tfrecord dirs: %s", tfrecord_dirs)

  for idx, record_dir in enumerate(tfrecord_dirs):
    record_glob = os.path.join(record_dir, record_glob_base)
    logging.info("[%d] Record glob: %s", idx, record_glob)

    record_paths = sorted(tf.io.gfile.glob(record_glob))
    logging.info("[%d] Num of record info path: %d", idx, len(record_paths))

    cur_record_info = {"num_batch": 0, "filenames": []}

    for record_info_path in record_paths:
      with tf.io.gfile.GFile(record_info_path, "r") as fp:
        info = json.load(fp)
        cur_record_info["num_batch"] += int(_get_num_batch(info))
        cur_record_info["filenames"] += info["filenames"]

    # overwrite directory for `cur_record_info`
    new_filenames = []
    for filename in cur_record_info["filenames"]:
      basename = os.path.basename(filename)
      new_filename = os.path.join(record_dir, basename)
      new_filenames.append(new_filename)
    cur_record_info["filenames"] = new_filenames

    logging.info("[Dir %d] Number of chosen batches: %s", idx,
                 cur_record_info["num_batch"])
    logging.info("[Dir %d] Number of chosen files: %s", idx,
                 len(cur_record_info["filenames"]))
    logging.info(cur_record_info["filenames"])

    # add `cur_record_info` to global `record_info`
    record_info["num_batch"] += cur_record_info["num_batch"]
    record_info["filenames"] += cur_record_info["filenames"]

  logging.info("Total number of batches: %d", record_info["num_batch"])
  logging.info("Total number of files: %d", len(record_info["filenames"]))
  logging.info(record_info["filenames"])

  def _dataset_fn(ctx=None):
    """Function that can create a pretrain dataset."""

    train_dataset = create_pretrain_dataset(
        file_names=record_info["filenames"],
        bsz_per_core=batch_size,
        seq_len=seq_len,
        reuse_len=reuse_len,
        perm_size=perm_size,
        leak_ratio=leak_ratio,
        online_masking_config=online_masking_config,
        num_predict=num_predict,
        input_pipeline_context=ctx)
    return train_dataset

  return _dataset_fn if use_dataset_fn else _dataset_fn()


def parse_files_to_dataset(parser,
                           file_paths,
                           bsz_per_core,
                           sequential,
                           input_pipeline_context=None):
  """Creates the dataset given file paths."""

  dataset = tf.data.Dataset.from_tensor_slices(file_paths)

  # Note: we cannot perform sample-level shuffle here because this will violate
  # the consecutive requirement of data stream.

  if input_pipeline_context and input_pipeline_context.num_input_pipelines > 1:
    dataset = dataset.shard(input_pipeline_context.num_input_pipelines,
                            input_pipeline_context.input_pipeline_id)
  # file-level shuffle
  if len(file_paths) > 1:
    dataset = dataset.shuffle(len(file_paths))

  if sequential:
    # Note: cannot perform sample-level shuffle here because this will violate
    # the consecutive requirement of data stream.
    dataset = tf.data.TFRecordDataset(dataset)
  else:
    # `cycle_length` is the number of parallel files that get read.
    cycle_length = min(8, len(file_paths))
    logging.info("Interleave %d files", cycle_length)

    dataset = dataset.apply(
        tf.data.experimental.parallel_interleave(
            tf.data.TFRecordDataset, cycle_length=cycle_length))
    buffer_size = 2048
    logging.info("Perform sample-level shuffle with size %d", buffer_size)
    dataset = dataset.shuffle(buffer_size=buffer_size)

  dataset = dataset.cache().repeat().map(parser)
  dataset = dataset.batch(bsz_per_core, drop_remainder=True)
  dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)

  return dataset


def _local_perm(inputs, is_masked, perm_size, seq_len, leak_ratio):
  """Samples a permutation of the factorization order.

     Creates perm_mask and target_mask accordingly.

  Args:
    inputs: int64 Tensor in shape [seq_len], input ids.
    is_masked: bool Tensor in shape [seq_len]. True means being selected for
      partial prediction.
    perm_size: the length of longest permutation. Could be set to be reuse_len.
      Should not be larger than reuse_len or there will be data leaks.
    seq_len: int, sequence length.
    leak_ratio: float, percent of masked tokens that are leaked.

  Returns:
    perm_mask: float32 Tensor in shape [seq_len, seq_len] consisted of 0 and 1.
    If perm_mask[i][j] == 1, it means the ith token (in original order) cannot
    attend to the jth token
    (in original order). This case will happen only when the ith token's
    permutated position <= the jth token's permutated position,
    and the jth token is masked or is func token. If perm_mask[i][j] == 0, it
    means the ith token (in original order) can attend to the jth token
    (in original order). Note that non-masked tokens can be attended by all
    other tokens, which is different from the description in original paper.
    target_mask: float32 Tensor in shape [seq_len] consisted of 0 and 1. If
    target_mask[i] == 1,
    the ith token needs to be predicted and mask will be used as input. This
    token will count for loss.
    If target_mask[i] == 0, token (or [SEP], [CLS]) will be used as input. This
    token will not count for loss.
    inputs_k: int64 Tensor in shape [seq_len], input ids.
    inputs_q: float32 Tensor in shape [seq_len], the same as target_mask.

  """

  # Generate permutation indices
  index = tf.range(seq_len, dtype=tf.int64)
  index = tf.transpose(tf.reshape(index, [-1, perm_size]))
  index = tf.random.shuffle(index)
  index = tf.reshape(tf.transpose(index), [-1])

  # non-functional tokens
  non_func_tokens = tf.logical_not(
      tf.logical_or(tf.equal(inputs, SEP_ID), tf.equal(inputs, CLS_ID)))
  masked_tokens = tf.logical_and(is_masked, non_func_tokens)
  non_masked_or_func_tokens = tf.logical_not(masked_tokens)

  smallest_index = -2 * tf.ones([seq_len], dtype=tf.int64)

  # Similar to BERT, randomly leak some masked tokens
  if leak_ratio > 0:
    leak_tokens = tf.logical_and(
        masked_tokens,
        tf.random.uniform([seq_len], maxval=1.0) < leak_ratio)
    can_attend_self = tf.logical_or(non_masked_or_func_tokens, leak_tokens)
  else:
    can_attend_self = non_masked_or_func_tokens
  to_index = tf.where(can_attend_self, smallest_index, index)
  from_index = tf.where(can_attend_self, to_index + 1, to_index)

  # For masked tokens, can attend if i > j
  # For context tokens, always can attend each other
  can_attend = from_index[:, None] > to_index[None, :]

  # In modeling, 1 indicates cannot attend. Hence, reverse the value here.
  perm_mask = 1.0 - tf.cast(can_attend, tf.float32)

  # Only masked tokens are included in the loss
  target_mask = tf.cast(masked_tokens, tf.float32)

  # construct inputs_k
  inputs_k = inputs

  # construct inputs_q
  inputs_q = masked_tokens

  return perm_mask, target_mask, inputs_k, inputs_q