from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
from bert import modeling
import collections
import tensorflow as tf

flags = tf.flags

FLAGS = flags.FLAGS

## Required parameters
flags.DEFINE_string(
    "bert_config_file", None,
    "The config json file corresponding to the pre-trained BERT model. "
    "This specifies the model architecture.")

flags.DEFINE_string(
    "input_file", None,
    "Input TF example files (can be a glob or comma separated).")

flags.DEFINE_string(
    "output_dir", None,
    "The output directory where the model checkpoints will be written.")

flags.DEFINE_integer("truncation_factor", 128,
                     "Number of top probable words to save from teacher network output")

## Other parameters
flags.DEFINE_string(
    "init_checkpoint", r'D:\Awake\Code\DistillBERT\bert-base-chinese\bert_model.ckpt',
    "Initial checkpoint (usually from a pre-trained BERT model).")

flags.DEFINE_integer(
    "max_seq_length", 512,
    "The maximum total input sequence length after WordPiece tokenization. "
    "Sequences longer than this will be truncated, and sequences shorter "
    "than this will be padded. Must match data generation.")

flags.DEFINE_integer(
    "max_predictions_per_seq", 20,
    "Maximum number of masked LM predictions per sequence. "
    "Must match data generation.")

flags.DEFINE_integer("predict_batch_size", 4, "Total batch size.")

flags.DEFINE_bool("tpu", False, "Whether to use TPU or GPU/CPU.")

tf.flags.DEFINE_string("master", None, "[Optional] TensorFlow master URL.")

tf.flags.DEFINE_string(
    "tpu_name", None,
    "The Cloud TPU to use for training. This should be either the name "
    "used when creating the Cloud TPU, or a grpc://ip.address.of.tpu:8470 "
    "url.")

tf.flags.DEFINE_string(
    "tpu_zone", None,
    "[Optional] GCE zone where the Cloud TPU is located in. If not "
    "specified, we will attempt to automatically detect the GCE project from "
    "metadata.")

tf.flags.DEFINE_string(
    "gcp_project", None,
    "[Optional] Project name for the Cloud TPU-enabled project. If not "
    "specified, we will attempt to automatically detect the GCE project from "
    "metadata.")

flags.DEFINE_integer(
    "num_tpu_cores", 8,
    "Only used if `use_tpu` is True. Total number of TPU cores to use.")

def input_fn_builder(input_files):
    """Decodes a record to a TensorFlow example."""
    name_to_features = {
        "input_ids":
            tf.FixedLenFeature([FLAGS.max_seq_length], tf.int64),
        "input_mask":
            tf.FixedLenFeature([FLAGS.max_seq_length], tf.int64),
        "segment_ids":
            tf.FixedLenFeature([FLAGS.max_seq_length], tf.int64),
        "masked_lm_positions":
            tf.FixedLenFeature([FLAGS.max_predictions_per_seq], tf.int64),
        "masked_lm_ids":
            tf.FixedLenFeature([FLAGS.max_predictions_per_seq], tf.int64),
        "masked_lm_weights":
            tf.FixedLenFeature([FLAGS.max_predictions_per_seq], tf.float32),
        "next_sentence_labels":
            tf.FixedLenFeature([1], tf.int64),
    }

    def _decode_record(record, name_to_features):
        example = tf.parse_single_example(record, name_to_features)

        for name in list(example.keys()):
            t = example[name]
            if t.dtype == tf.int64:
                t = tf.to_int32(t)
            example[name] = t

        return example

    def input_fn(params):
        batch_size = params["batch_size"]
        d = tf.data.TFRecordDataset(input_files)
        # d = tf.data.Dataset.from_tensor_slices(tf.constant(input_files))
        d = d.repeat(1)

        d = d.apply(
            tf.contrib.data.map_and_batch(
                lambda record: _decode_record(record, name_to_features),
                batch_size=batch_size,
                drop_remainder=False))

        return d

    return input_fn

def get_masked_lm_output(bert_config, input_tensor, output_weights, positions, truncation_factor):
  """Get loss and log probs for the masked LM."""
  input_tensor = gather_indexes(input_tensor, positions)

  with tf.variable_scope("cls/predictions"):
    # We apply one more non-linear transformation before the output layer.
    # This matrix is not used after pre-training.
    with tf.variable_scope("transform"):
      input_tensor = tf.layers.dense(
          input_tensor,
          units=bert_config.hidden_size,
          activation=modeling.get_activation(bert_config.hidden_act),
          kernel_initializer=modeling.create_initializer(
              bert_config.initializer_range))
      input_tensor = modeling.layer_norm(input_tensor)

    # The output weights are the same as the input embeddings, but there is
    # an output-only bias for each token.
    output_bias = tf.get_variable(
        "output_bias",
        shape=[bert_config.vocab_size],
        initializer=tf.zeros_initializer())
    logits = tf.matmul(input_tensor, output_weights, transpose_b=True)
    logits = tf.nn.bias_add(logits, output_bias)
    masked_lm_probs = tf.nn.softmax(logits, axis=-1)
    trunc_masked_lm_probs, top_indices = tf.math.top_k(masked_lm_probs, k=truncation_factor, sorted=False)

    max_predictions_per_seq = positions.get_shape().as_list()[1]
    truncation_factor_ = top_indices.get_shape().as_list()[1]

    trunc_masked_lm_probs = tf.reshape(trunc_masked_lm_probs, [-1, max_predictions_per_seq, truncation_factor_])
    top_indices = tf.reshape(top_indices, [-1, max_predictions_per_seq, truncation_factor_])
  return trunc_masked_lm_probs, top_indices

def get_next_sentence_output(bert_config, input_tensor):
  """Get loss and log probs for the next sentence prediction."""

  # Simple binary classification. Note that 0 is "next sentence" and 1 is
  # "random sentence". This weight matrix is not used after pre-training.
  with tf.variable_scope("cls/seq_relationship"):
    output_weights = tf.get_variable(
        "output_weights",
        shape=[2, bert_config.hidden_size],
        initializer=modeling.create_initializer(bert_config.initializer_range))
    output_bias = tf.get_variable(
        "output_bias", shape=[2], initializer=tf.zeros_initializer())

    logits = tf.matmul(input_tensor, output_weights, transpose_b=True)
    logits = tf.nn.bias_add(logits, output_bias)
    next_sentence_probs = tf.nn.softmax(logits, axis=-1)

    return next_sentence_probs

def gather_indexes(sequence_tensor, positions):
  """Gathers the vectors at the specific positions over a minibatch."""
  sequence_shape = modeling.get_shape_list(sequence_tensor, expected_rank=3)
  batch_size = sequence_shape[0]
  seq_length = sequence_shape[1]
  width = sequence_shape[2]

  flat_offsets = tf.reshape(
      tf.range(0, batch_size, dtype=tf.int32) * seq_length, [-1, 1])
  flat_positions = tf.reshape(positions + flat_offsets, [-1])
  flat_sequence_tensor = tf.reshape(sequence_tensor,
                                    [batch_size * seq_length, width])
  output_tensor = tf.gather(flat_sequence_tensor, flat_positions)

  return output_tensor


def model_fn_builder(bert_config, init_checkpoint, truncation_factor):

    """The `model_fn` for TPUEstimator."""
    def model_fn(features, labels, mode, params):
        tf.logging.info("*** Features ***")
        for name in sorted(features.keys()):
          tf.logging.info("  name = %s, shape = %s" % (name, features[name].shape))

        input_ids = features["input_ids"]
        input_mask = features["input_mask"]
        segment_ids = features["segment_ids"]
        masked_lm_positions = features["masked_lm_positions"]

        model = modeling.BertModel(
            config=bert_config,
            is_training=False,
            input_ids=input_ids,
            input_mask=input_mask,
            token_type_ids=segment_ids,
            use_one_hot_embeddings=False)

        trunc_masked_lm_probs, top_indices = get_masked_lm_output(bert_config,
                                                   model.get_sequence_output(),
                                                   model.get_embedding_table(),
                                                   masked_lm_positions,
                                                   truncation_factor)

        tvars = tf.trainable_variables()

        initialized_variable_names = {}
        scaffold_fn = None

        if init_checkpoint:
            (assignment_map, initialized_variable_names) = modeling.get_assignment_map_from_checkpoint(tvars, init_checkpoint)
            tf.train.init_from_checkpoint(init_checkpoint, assignment_map)

            if FLAGS.tpu:
                def tpu_scaffold():
                    tf.train.init_from_checkpoint(init_checkpoint, assignment_map)
                    return tf.train.Scaffold()

                scaffold_fn = tpu_scaffold

            else:
                tf.train.init_from_checkpoint(init_checkpoint, assignment_map)

        tf.logging.info("**** Trainable Variables ****")
        for var in tvars:
            init_string = ""
            if var.name in initialized_variable_names:
                init_string = ", *INIT_FROM_CKPT*"
            tf.logging.info("  name = %s, shape = %s%s", var.name, var.shape, init_string)

        output_spec = None
        if mode == tf.estimator.ModeKeys.PREDICT:
            features['masked_lm_probs'] = trunc_masked_lm_probs
            features['top_indices'] = top_indices

            output_spec = tf.contrib.tpu.TPUEstimatorSpec(mode=mode, predictions=features, scaffold_fn=scaffold_fn)

        else:
            raise ValueError('this is only for predict, please change the estimator to predict')

        return output_spec

    return model_fn

def create_int_feature(values):
  feature = tf.train.Feature(int64_list=tf.train.Int64List(value=list(values)))
  return feature


def create_float_feature(values):
  feature = tf.train.Feature(float_list=tf.train.FloatList(value=list(values)))
  return feature


def main(_):
    tf.logging.set_verbosity(tf.logging.INFO)

    input_files = []

    tf.gfile.MakeDirs(FLAGS.output_dir)

    for input_pattern in FLAGS.input_file.split(","):
        input_files.extend(tf.gfile.Glob(input_pattern))

    tf.logging.info("*** Input Files ***")
    for input_file in input_files:
        tf.logging.info("  %s" % input_file)

    tf.contrib.tpu.RunConfig()

    tpu_cluster_resolver = None
    if FLAGS.tpu and FLAGS.tpu_name:
        tpu_cluster_resolver = tf.contrib.cluster_resolver.TPUClusterResolver(
            FLAGS.tpu_name, zone=FLAGS.tpu_zone, project=FLAGS.gcp_project)

    is_per_host = tf.contrib.tpu.InputPipelineConfig.PER_HOST_V2

    run_config = tf.contrib.tpu.RunConfig(
        cluster=tpu_cluster_resolver,
        master=FLAGS.master,
        tpu_config=tf.contrib.tpu.TPUConfig(
            num_shards=FLAGS.num_tpu_cores,
            per_host_input_for_training=is_per_host))

    bert_config = modeling.BertConfig.from_json_file(FLAGS.bert_config_file)
    model_fn = model_fn_builder(bert_config, FLAGS.init_checkpoint, FLAGS.truncation_factor)

    estimator = tf.contrib.tpu.TPUEstimator(
            use_tpu=FLAGS.tpu,
            model_fn=model_fn,
            config=run_config,
            train_batch_size=256,
            predict_batch_size=FLAGS.predict_batch_size)

    input_fn = input_fn_builder(input_files)

    writer = None
    count = 0
    tf.logging.info('start create teacher output!')
    for features in estimator.predict(input_fn, yield_single_examples=False):
        masked_lm_probs = features['masked_lm_probs']
        top_indices = features['top_indices']

        input_ids = features["input_ids"]
        input_mask = features["input_mask"]
        segment_ids = features["segment_ids"]
        masked_lm_positions = features["masked_lm_positions"]
        masked_lm_ids = features["masked_lm_ids"]
        masked_lm_weights = features["masked_lm_weights"]
        next_sentence_labels = features["next_sentence_labels"]

        cur_batch_size, seq_len = input_ids.shape
        _, pred_per_seq = masked_lm_positions.shape
        feature = collections.OrderedDict()
        for i in range(cur_batch_size):
            if count % 50000 == 0:
                writer = tf.python_io.TFRecordWriter(FLAGS.output_dir + 'teacher_%d.tfrecord' % (count / 50000))

            feature['input_ids'] = create_int_feature(input_ids[i])
            feature["input_mask"] = create_int_feature(input_mask[i])
            feature["segment_ids"] = create_int_feature(segment_ids[i])
            feature["masked_lm_positions"] = create_int_feature(masked_lm_positions[i])
            feature["masked_lm_ids"] = create_int_feature(masked_lm_ids[i])
            feature["masked_lm_weights"] = create_float_feature(masked_lm_weights[i])
            feature["next_sentence_labels"] = create_int_feature([next_sentence_labels[i][0]])

            # retrieve predictions for (pred_per_seq) number of words, and flatten
            distribution = masked_lm_probs[i].reshape(-1)
            feature["truncated_masked_lm_probs"] = create_float_feature(distribution)

            top_k_dex = top_indices[i].reshape(-1)
            feature["top_k_indices"] = create_int_feature(top_k_dex)

            tf_example = tf.train.Example(features=tf.train.Features(feature=feature))
            writer.write(tf_example.SerializeToString())
            count += 1

    writer.close()
    tf.logging.info('create teacher output over!')

if __name__ =="__main__":
    tf.app.run()
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
from bert import modeling
import collections
import tensorflow as tf

flags = tf.flags

FLAGS = flags.FLAGS

## Required parameters
flags.DEFINE_string(
    "bert_config_file", None,
    "The config json file corresponding to the pre-trained BERT model. "
    "This specifies the model architecture.")

flags.DEFINE_string(
    "input_file", None,
    "Input TF example files (can be a glob or comma separated).")

flags.DEFINE_string(
    "output_dir", None,
    "The output directory where the model checkpoints will be written.")

flags.DEFINE_integer("truncation_factor", 128,
                     "Number of top probable words to save from teacher network output")

## Other parameters
flags.DEFINE_string(
    "init_checkpoint", r'D:\Awake\Code\DistillBERT\bert-base-chinese\bert_model.ckpt',
    "Initial checkpoint (usually from a pre-trained BERT model).")

flags.DEFINE_integer(
    "max_seq_length", 512,
    "The maximum total input sequence length after WordPiece tokenization. "
    "Sequences longer than this will be truncated, and sequences shorter "
    "than this will be padded. Must match data generation.")

flags.DEFINE_integer(
    "max_predictions_per_seq", 20,
    "Maximum number of masked LM predictions per sequence. "
    "Must match data generation.")

flags.DEFINE_integer("predict_batch_size", 4, "Total batch size.")

flags.DEFINE_bool("tpu", False, "Whether to use TPU or GPU/CPU.")

tf.flags.DEFINE_string("master", None, "[Optional] TensorFlow master URL.")

tf.flags.DEFINE_string(
    "tpu_name", None,
    "The Cloud TPU to use for training. This should be either the name "
    "used when creating the Cloud TPU, or a grpc://ip.address.of.tpu:8470 "
    "url.")

tf.flags.DEFINE_string(
    "tpu_zone", None,
    "[Optional] GCE zone where the Cloud TPU is located in. If not "
    "specified, we will attempt to automatically detect the GCE project from "
    "metadata.")

tf.flags.DEFINE_string(
    "gcp_project", None,
    "[Optional] Project name for the Cloud TPU-enabled project. If not "
    "specified, we will attempt to automatically detect the GCE project from "
    "metadata.")

flags.DEFINE_integer(
    "num_tpu_cores", 8,
    "Only used if `use_tpu` is True. Total number of TPU cores to use.")

def input_fn_builder(input_files):
    """Decodes a record to a TensorFlow example."""
    name_to_features = {
        "input_ids":
            tf.FixedLenFeature([FLAGS.max_seq_length], tf.int64),
        "input_mask":
            tf.FixedLenFeature([FLAGS.max_seq_length], tf.int64),
        "segment_ids":
            tf.FixedLenFeature([FLAGS.max_seq_length], tf.int64),
        "masked_lm_positions":
            tf.FixedLenFeature([FLAGS.max_predictions_per_seq], tf.int64),
        "masked_lm_ids":
            tf.FixedLenFeature([FLAGS.max_predictions_per_seq], tf.int64),
        "masked_lm_weights":
            tf.FixedLenFeature([FLAGS.max_predictions_per_seq], tf.float32),
        "next_sentence_labels":
            tf.FixedLenFeature([1], tf.int64),
    }

    def _decode_record(record, name_to_features):
        example = tf.parse_single_example(record, name_to_features)

        for name in list(example.keys()):
            t = example[name]
            if t.dtype == tf.int64:
                t = tf.to_int32(t)
            example[name] = t

        return example

    def input_fn(params):
        batch_size = params["batch_size"]
        d = tf.data.TFRecordDataset(input_files)
        # d = tf.data.Dataset.from_tensor_slices(tf.constant(input_files))
        d = d.repeat(1)

        d = d.apply(
            tf.contrib.data.map_and_batch(
                lambda record: _decode_record(record, name_to_features),
                batch_size=batch_size,
                drop_remainder=False))

        return d

    return input_fn

def get_masked_lm_output(bert_config, input_tensor, output_weights, positions, truncation_factor):
  """Get loss and log probs for the masked LM."""
  input_tensor = gather_indexes(input_tensor, positions)

  with tf.variable_scope("cls/predictions"):
    # We apply one more non-linear transformation before the output layer.
    # This matrix is not used after pre-training.
    with tf.variable_scope("transform"):
      input_tensor = tf.layers.dense(
          input_tensor,
          units=bert_config.hidden_size,
          activation=modeling.get_activation(bert_config.hidden_act),
          kernel_initializer=modeling.create_initializer(
              bert_config.initializer_range))
      input_tensor = modeling.layer_norm(input_tensor)

    # The output weights are the same as the input embeddings, but there is
    # an output-only bias for each token.
    output_bias = tf.get_variable(
        "output_bias",
        shape=[bert_config.vocab_size],
        initializer=tf.zeros_initializer())
    logits = tf.matmul(input_tensor, output_weights, transpose_b=True)
    logits = tf.nn.bias_add(logits, output_bias)
    masked_lm_probs = tf.nn.softmax(logits, axis=-1)
    trunc_masked_lm_probs, top_indices = tf.math.top_k(masked_lm_probs, k=truncation_factor, sorted=False)

    max_predictions_per_seq = positions.get_shape().as_list()[1]
    truncation_factor_ = top_indices.get_shape().as_list()[1]

    trunc_masked_lm_probs = tf.reshape(trunc_masked_lm_probs, [-1, max_predictions_per_seq, truncation_factor_])
    top_indices = tf.reshape(top_indices, [-1, max_predictions_per_seq, truncation_factor_])
  return trunc_masked_lm_probs, top_indices

def get_next_sentence_output(bert_config, input_tensor):
  """Get loss and log probs for the next sentence prediction."""

  # Simple binary classification. Note that 0 is "next sentence" and 1 is
  # "random sentence". This weight matrix is not used after pre-training.
  with tf.variable_scope("cls/seq_relationship"):
    output_weights = tf.get_variable(
        "output_weights",
        shape=[2, bert_config.hidden_size],
        initializer=modeling.create_initializer(bert_config.initializer_range))
    output_bias = tf.get_variable(
        "output_bias", shape=[2], initializer=tf.zeros_initializer())

    logits = tf.matmul(input_tensor, output_weights, transpose_b=True)
    logits = tf.nn.bias_add(logits, output_bias)
    next_sentence_probs = tf.nn.softmax(logits, axis=-1)

    return next_sentence_probs

def gather_indexes(sequence_tensor, positions):
  """Gathers the vectors at the specific positions over a minibatch."""
  sequence_shape = modeling.get_shape_list(sequence_tensor, expected_rank=3)
  batch_size = sequence_shape[0]
  seq_length = sequence_shape[1]
  width = sequence_shape[2]

  flat_offsets = tf.reshape(
      tf.range(0, batch_size, dtype=tf.int32) * seq_length, [-1, 1])
  flat_positions = tf.reshape(positions + flat_offsets, [-1])
  flat_sequence_tensor = tf.reshape(sequence_tensor,
                                    [batch_size * seq_length, width])
  output_tensor = tf.gather(flat_sequence_tensor, flat_positions)

  return output_tensor


def model_fn_builder(bert_config, init_checkpoint, truncation_factor):

    """The `model_fn` for TPUEstimator."""
    def model_fn(features, labels, mode, params):
        tf.logging.info("*** Features ***")
        for name in sorted(features.keys()):
          tf.logging.info("  name = %s, shape = %s" % (name, features[name].shape))

        input_ids = features["input_ids"]
        input_mask = features["input_mask"]
        segment_ids = features["segment_ids"]
        masked_lm_positions = features["masked_lm_positions"]

        model = modeling.BertModel(
            config=bert_config,
            is_training=False,
            input_ids=input_ids,
            input_mask=input_mask,
            token_type_ids=segment_ids,
            use_one_hot_embeddings=False)

        trunc_masked_lm_probs, top_indices = get_masked_lm_output(bert_config,
                                                   model.get_sequence_output(),
                                                   model.get_embedding_table(),
                                                   masked_lm_positions,
                                                   truncation_factor)

        tvars = tf.trainable_variables()

        initialized_variable_names = {}
        scaffold_fn = None

        if init_checkpoint:
            (assignment_map, initialized_variable_names) = modeling.get_assignment_map_from_checkpoint(tvars, init_checkpoint)
            tf.train.init_from_checkpoint(init_checkpoint, assignment_map)

            if FLAGS.tpu:
                def tpu_scaffold():
                    tf.train.init_from_checkpoint(init_checkpoint, assignment_map)
                    return tf.train.Scaffold()

                scaffold_fn = tpu_scaffold

            else:
                tf.train.init_from_checkpoint(init_checkpoint, assignment_map)

        tf.logging.info("**** Trainable Variables ****")
        for var in tvars:
            init_string = ""
            if var.name in initialized_variable_names:
                init_string = ", *INIT_FROM_CKPT*"
            tf.logging.info("  name = %s, shape = %s%s", var.name, var.shape, init_string)

        output_spec = None
        if mode == tf.estimator.ModeKeys.PREDICT:
            features['masked_lm_probs'] = trunc_masked_lm_probs
            features['top_indices'] = top_indices

            output_spec = tf.contrib.tpu.TPUEstimatorSpec(mode=mode, predictions=features, scaffold_fn=scaffold_fn)

        else:
            raise ValueError('this is only for predict, please change the estimator to predict')

        return output_spec

    return model_fn

def create_int_feature(values):
  feature = tf.train.Feature(int64_list=tf.train.Int64List(value=list(values)))
  return feature


def create_float_feature(values):
  feature = tf.train.Feature(float_list=tf.train.FloatList(value=list(values)))
  return feature


def main(_):
    tf.logging.set_verbosity(tf.logging.INFO)

    input_files = []

    tf.gfile.MakeDirs(FLAGS.output_dir)

    for input_pattern in FLAGS.input_file.split(","):
        input_files.extend(tf.gfile.Glob(input_pattern))

    tf.logging.info("*** Input Files ***")
    for input_file in input_files:
        tf.logging.info("  %s" % input_file)

    tf.contrib.tpu.RunConfig()

    tpu_cluster_resolver = None
    if FLAGS.tpu and FLAGS.tpu_name:
        tpu_cluster_resolver = tf.contrib.cluster_resolver.TPUClusterResolver(
            FLAGS.tpu_name, zone=FLAGS.tpu_zone, project=FLAGS.gcp_project)

    is_per_host = tf.contrib.tpu.InputPipelineConfig.PER_HOST_V2

    run_config = tf.contrib.tpu.RunConfig(
        cluster=tpu_cluster_resolver,
        master=FLAGS.master,
        tpu_config=tf.contrib.tpu.TPUConfig(
            num_shards=FLAGS.num_tpu_cores,
            per_host_input_for_training=is_per_host))

    bert_config = modeling.BertConfig.from_json_file(FLAGS.bert_config_file)
    model_fn = model_fn_builder(bert_config, FLAGS.init_checkpoint, FLAGS.truncation_factor)

    estimator = tf.contrib.tpu.TPUEstimator(
            use_tpu=FLAGS.tpu,
            model_fn=model_fn,
            config=run_config,
            train_batch_size=256,
            predict_batch_size=FLAGS.predict_batch_size)

    input_fn = input_fn_builder(input_files)

    writer = None
    count = 0
    tf.logging.info('start create teacher output!')
    for features in estimator.predict(input_fn, yield_single_examples=False):
        masked_lm_probs = features['masked_lm_probs']
        top_indices = features['top_indices']

        input_ids = features["input_ids"]
        input_mask = features["input_mask"]
        segment_ids = features["segment_ids"]
        masked_lm_positions = features["masked_lm_positions"]
        masked_lm_ids = features["masked_lm_ids"]
        masked_lm_weights = features["masked_lm_weights"]
        next_sentence_labels = features["next_sentence_labels"]

        cur_batch_size, seq_len = input_ids.shape
        _, pred_per_seq = masked_lm_positions.shape
        feature = collections.OrderedDict()
        for i in range(cur_batch_size):
            if count % 50000 == 0:
                writer = tf.python_io.TFRecordWriter(FLAGS.output_dir + 'teacher_%d.tfrecord' % (count / 50000))

            feature['input_ids'] = create_int_feature(input_ids[i])
            feature["input_mask"] = create_int_feature(input_mask[i])
            feature["segment_ids"] = create_int_feature(segment_ids[i])
            feature["masked_lm_positions"] = create_int_feature(masked_lm_positions[i])
            feature["masked_lm_ids"] = create_int_feature(masked_lm_ids[i])
            feature["masked_lm_weights"] = create_float_feature(masked_lm_weights[i])
            feature["next_sentence_labels"] = create_int_feature([next_sentence_labels[i][0]])

            # retrieve predictions for (pred_per_seq) number of words, and flatten
            distribution = masked_lm_probs[i].reshape(-1)
            feature["truncated_masked_lm_probs"] = create_float_feature(distribution)

            top_k_dex = top_indices[i].reshape(-1)
            feature["top_k_indices"] = create_int_feature(top_k_dex)

            tf_example = tf.train.Example(features=tf.train.Features(feature=feature))
            writer.write(tf_example.SerializeToString())
            count += 1

    writer.close()
    tf.logging.info('create teacher output over!')

if __name__ =="__main__":
    tf.app.run()
## Required parameters ## Other parameters Decodes a record to a TensorFlow example. # d = tf.data.Dataset.from_tensor_slices(tf.constant(input_files)) Get loss and log probs for the masked LM. # We apply one more non-linear transformation before the output layer. # This matrix is not used after pre-training. # The output weights are the same as the input embeddings, but there is # an output-only bias for each token. Get loss and log probs for the next sentence prediction. # Simple binary classification. Note that 0 is "next sentence" and 1 is # "random sentence". This weight matrix is not used after pre-training. Gathers the vectors at the specific positions over a minibatch. The `model_fn` for TPUEstimator. # retrieve predictions for (pred_per_seq) number of words, and flatten