# coding=utf-8
# Copyright (c) 2019 NVIDIA CORPORATION. All rights reserved.
# Copyright 2018 The Google AI Language Team Authors.
#
# 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.

"""BERT finetuning runner."""

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

import collections
import csv
import os
import modeling
import optimization
import tokenization
import tensorflow as tf
import horovod.tensorflow as hvd
import time
from utils.utils import LogEvalRunHook, LogTrainRunHook, setup_xla_flags
from utils.gpu_affinity import set_affinity
import utils.dllogger_class
from dllogger import Verbosity
from utils.create_glue_data import *
import numpy as np
import tf_metrics

flags = tf.flags

FLAGS = flags.FLAGS

## Required parameters
flags.DEFINE_string(
    "data_dir", None,
    "The input data dir. Should contain the .tsv files (or other data files) "
    "for the task.")

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("task_name", None, "The name of the task to train.")

flags.DEFINE_string("vocab_file", None,
                    "The vocabulary file that the BERT model was trained on.")

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

## Other parameters
flags.DEFINE_string(
    "dllog_path", "/results/bert_dllog.json",
    "filename where dllogger writes to")

flags.DEFINE_string(
    "optimizer_type", "lamb",
    "Optimizer type : adam or lamb")

flags.DEFINE_string(
    "init_checkpoint", None,
    "Initial checkpoint (usually from a pre-trained BERT model).")

flags.DEFINE_bool(
    "do_lower_case", True,
    "Whether to lower case the input text. Should be True for uncased "
    "models and False for cased models.")

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

flags.DEFINE_bool("do_train", False, "Whether to run training.")

flags.DEFINE_bool("do_eval", False, "Whether to run eval on the dev set.")

flags.DEFINE_bool(
    "do_predict", False,
    "Whether to run the model in inference mode on the test set.")

flags.DEFINE_integer("train_batch_size", 32, "Total batch size for training.")

flags.DEFINE_integer("eval_batch_size", 8, "Total batch size for eval.")

flags.DEFINE_integer("predict_batch_size", 8, "Total batch size for predict.")

flags.DEFINE_float("learning_rate", 5e-5, "The initial learning rate for Adam.")

flags.DEFINE_bool("use_trt", False, "Whether to use TF-TRT")

flags.DEFINE_float("num_train_epochs", 3.0,
                   "Total number of training epochs to perform.")

flags.DEFINE_float(
    "warmup_proportion", 0.1,
    "Proportion of training to perform linear learning rate warmup for. "
    "E.g., 0.1 = 10% of training.")

flags.DEFINE_integer("save_checkpoints_steps", 1000,
                     "How often to save the model checkpoint.")
flags.DEFINE_integer("display_loss_steps", 10,
                     "How often to print loss from estimator")

flags.DEFINE_integer("iterations_per_loop", 1000,
                     "How many steps to make in each estimator call.")
flags.DEFINE_integer("num_accumulation_steps", 1,
                     "Number of accumulation steps before gradient update" 
                      "Global batch size = num_accumulation_steps * train_batch_size")
flags.DEFINE_bool("amp", True, "Whether to enable AMP ops. When false, uses TF32 on A100 and FP32 on V100 GPUS.")
flags.DEFINE_bool("use_xla", True, "Whether to enable XLA JIT compilation.")
flags.DEFINE_bool("horovod", False, "Whether to use Horovod for multi-gpu runs")

flags.DEFINE_bool(
    "verbose_logging", False,
    "If true, all of the warnings related to data processing will be printed. "
    "A number of warnings are expected for a normal SQuAD evaluation.")


def file_based_input_fn_builder(input_file, batch_size, seq_length, is_training,
                                drop_remainder, hvd=None):
  """Creates an `input_fn` closure to be passed to Estimator."""

  name_to_features = {
      "input_ids": tf.io.FixedLenFeature([seq_length], tf.int64),
      "input_mask": tf.io.FixedLenFeature([seq_length], tf.int64),
      "segment_ids": tf.io.FixedLenFeature([seq_length], tf.int64),
      "label_ids": tf.io.FixedLenFeature([], tf.int64),
  }

  def _decode_record(record, name_to_features):
    """Decodes a record to a TensorFlow example."""
    example = tf.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.to_int32(t)
      example[name] = t

    return example

  def input_fn():
    """The actual input function."""

    # For training, we want a lot of parallel reading and shuffling.
    # For eval, we want no shuffling and parallel reading doesn't matter.
    d = tf.data.TFRecordDataset(input_file)
    if is_training:
      if hvd is not None: d = d.shard(hvd.size(), hvd.rank())
      d = d.repeat()
      d = d.shuffle(buffer_size=100)

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

    return d

  return input_fn


def create_model(bert_config, is_training, input_ids, input_mask, segment_ids,
                 labels, num_labels, use_one_hot_embeddings):
  """Creates a classification model."""
  model = modeling.BertModel(
      config=bert_config,
      is_training=is_training,
      input_ids=input_ids,
      input_mask=input_mask,
      token_type_ids=segment_ids,
      use_one_hot_embeddings=use_one_hot_embeddings,
      compute_type=tf.float32)

  # In the demo, we are doing a simple classification task on the entire
  # segment.
  #
  # If you want to use the token-level output, use model.get_sequence_output()
  # instead.
  output_layer = model.get_pooled_output()

  hidden_size = output_layer.shape[-1].value

  output_weights = tf.get_variable(
      "output_weights", [num_labels, hidden_size],
      initializer=tf.truncated_normal_initializer(stddev=0.02))

  output_bias = tf.get_variable(
      "output_bias", [num_labels], initializer=tf.zeros_initializer())

  with tf.variable_scope("loss"):
    if is_training:
      # I.e., 0.1 dropout
      output_layer = tf.nn.dropout(output_layer, keep_prob=0.9)

    logits = tf.matmul(output_layer, output_weights, transpose_b=True)
    logits = tf.nn.bias_add(logits, output_bias, name='cls_logits')
    probabilities = tf.nn.softmax(logits, axis=-1, name='cls_probabilities')
    log_probs = tf.nn.log_softmax(logits, axis=-1)

    one_hot_labels = tf.one_hot(labels, depth=num_labels, dtype=tf.float32)

    per_example_loss = -tf.reduce_sum(one_hot_labels * log_probs, axis=-1, name='cls_per_example_loss')
    loss = tf.reduce_mean(per_example_loss, name='cls_loss')

    return (loss, per_example_loss, logits, probabilities)

def get_frozen_tftrt_model(bert_config, shape, num_labels, use_one_hot_embeddings, init_checkpoint):
  tf_config = tf.compat.v1.ConfigProto()
  tf_config.gpu_options.allow_growth = True
  output_node_names = ['loss/cls_loss', 'loss/cls_per_example_loss', 'loss/cls_logits', 'loss/cls_probabilities']

  with tf.Session(config=tf_config) as tf_sess:
    input_ids = tf.placeholder(tf.int32, shape, 'input_ids')
    input_mask = tf.placeholder(tf.int32, shape, 'input_mask')
    segment_ids = tf.placeholder(tf.int32, shape, 'segment_ids')
    label_ids = tf.placeholder(tf.int32, (None), 'label_ids')

    create_model(bert_config, False, input_ids, input_mask, segment_ids, label_ids,
            num_labels, use_one_hot_embeddings)

    tvars = tf.trainable_variables()
    (assignment_map, initialized_variable_names) = modeling.get_assignment_map_from_checkpoint(tvars, init_checkpoint)
    tf.train.init_from_checkpoint(init_checkpoint, assignment_map)
    tf_sess.run(tf.global_variables_initializer())
    print("LOADED!")
    tf.compat.v1.logging.info("**** Trainable Variables ****")
    for var in tvars:
      init_string = ""
      if var.name in initialized_variable_names:
        init_string = ", *INIT_FROM_CKPT*"
      else:
        init_string = ", *NOTTTTTTTTTTTTTTTTTTTTT"
        tf.compat.v1.logging.info("  name = %s, shape = %s%s", var.name, var.shape, init_string)

    frozen_graph = tf.graph_util.convert_variables_to_constants(tf_sess, 
            tf_sess.graph.as_graph_def(), output_node_names)

    num_nodes = len(frozen_graph.node)
    print('Converting graph using TensorFlow-TensorRT...')
    from tensorflow.python.compiler.tensorrt import trt_convert as trt
    converter = trt.TrtGraphConverter(
        input_graph_def=frozen_graph,
        nodes_blacklist=output_node_names,
        max_workspace_size_bytes=(4096 << 20) - 1000,
        precision_mode = "FP16" if FLAGS.amp else "FP32",
        minimum_segment_size=4,
        is_dynamic_op=True,
        maximum_cached_engines=1000
    )
    frozen_graph = converter.convert()

    print('Total node count before and after TF-TRT conversion:',
          num_nodes, '->', len(frozen_graph.node))
    print('TRT node count:',
          len([1 for n in frozen_graph.node if str(n.op) == 'TRTEngineOp']))
    
    with tf.io.gfile.GFile("frozen_modelTRT.pb", "wb") as f:
      f.write(frozen_graph.SerializeToString())      
        
  return frozen_graph



def model_fn_builder(task_name, bert_config, num_labels, init_checkpoint, learning_rate,
                     num_train_steps, num_warmup_steps,
                     use_one_hot_embeddings, hvd=None):
  """Returns `model_fn` closure for Estimator."""

  def model_fn(features, labels, mode, params):  # pylint: disable=unused-argument
    """The `model_fn` for Estimator."""

    def metric_fn(per_example_loss, label_ids, logits):
        predictions = tf.argmax(logits, axis=-1, output_type=tf.int32)
        if task_name == "cola":
            FN, FN_op = tf.metrics.false_negatives(labels=label_ids, predictions=predictions)
            FP, FP_op = tf.metrics.false_positives(labels=label_ids, predictions=predictions)
            TP, TP_op = tf.metrics.true_positives(labels=label_ids, predictions=predictions)
            TN, TN_op = tf.metrics.true_negatives(labels=label_ids, predictions=predictions)

            MCC = (TP * TN - FP * FN) / ((TP + FP) * (TP + FN) * (TN + FP) * (TN + FN)) ** 0.5
            MCC_op = tf.group(FN_op, TN_op, TP_op, FP_op, tf.identity(MCC, name="MCC"))
            return {"MCC": (MCC, MCC_op)}
        elif task_name == "mrpc":
            accuracy = tf.metrics.accuracy(
                labels=label_ids, predictions=predictions)
            loss = tf.metrics.mean(values=per_example_loss)
            f1 = tf_metrics.f1(labels=label_ids, predictions=predictions, num_classes=2, pos_indices=[1])
            return {
                "eval_accuracy": accuracy,
                "eval_f1": f1,
                "eval_loss": loss,
            }
        else:
            accuracy = tf.metrics.accuracy(
                labels=label_ids, predictions=predictions)
            loss = tf.metrics.mean(values=per_example_loss)
            return {
                "eval_accuracy": accuracy,
                "eval_loss": loss,
            }
    tf.compat.v1.logging.info("*** Features ***")
    tf.compat.v1.logging.info("*** Features ***")
    for name in sorted(features.keys()):
      tf.compat.v1.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"]
    label_ids = features["label_ids"]

    is_training = (mode == tf.estimator.ModeKeys.TRAIN)

    if not is_training and FLAGS.use_trt:
        trt_graph = get_frozen_tftrt_model(bert_config, input_ids.shape, num_labels, use_one_hot_embeddings, init_checkpoint)
        (total_loss, per_example_loss, logits, probabilities)  = tf.import_graph_def(trt_graph,
                input_map={'input_ids':input_ids, 'input_mask':input_mask, 'segment_ids':segment_ids, 'label_ids':label_ids},
                return_elements=['loss/cls_loss:0', 'loss/cls_per_example_loss:0', 'loss/cls_logits:0', 'loss/cls_probabilities:0'],
                name='')
        if mode == tf.estimator.ModeKeys.PREDICT:
            predictions = {"probabilities": probabilities}
            output_spec = tf.estimator.EstimatorSpec(
                mode=mode, predictions=predictions)
        elif mode == tf.estimator.ModeKeys.EVAL:
            eval_metric_ops = metric_fn(per_example_loss, label_ids, logits)
            output_spec = tf.estimator.EstimatorSpec(
                mode=mode,
                loss=total_loss,
                eval_metric_ops=eval_metric_ops)
        return output_spec
    (total_loss, per_example_loss, logits, probabilities) = create_model(
        bert_config, is_training, input_ids, input_mask, segment_ids, label_ids,
        num_labels, use_one_hot_embeddings)

    tvars = tf.trainable_variables()
    initialized_variable_names = {}
    if init_checkpoint and (hvd is None or hvd.rank() == 0):
      (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.verbose_logging:
        tf.compat.v1.logging.info("**** Trainable Variables ****")
        for var in tvars:
          init_string = ""
          if var.name in initialized_variable_names:
            init_string = ", *INIT_FROM_CKPT*"
          tf.compat.v1.logging.info("  name = %s, shape = %s%s", var.name, var.shape,
                          init_string)

    output_spec = None
    if mode == tf.estimator.ModeKeys.TRAIN:

      train_op = optimization.create_optimizer(
          total_loss, learning_rate, num_train_steps, num_warmup_steps,
          hvd, False, FLAGS.amp, FLAGS.num_accumulation_steps, FLAGS.optimizer_type)
      output_spec = tf.estimator.EstimatorSpec(
          mode=mode,
          loss=total_loss,
          train_op=train_op)
    elif mode == tf.estimator.ModeKeys.EVAL:
      dummy_op = tf.no_op()
      # Need to call mixed precision graph rewrite if fp16 to enable graph rewrite
      if FLAGS.amp:
        loss_scaler = tf.train.experimental.FixedLossScale(1)
        dummy_op = tf.train.experimental.enable_mixed_precision_graph_rewrite(
            optimization.LAMBOptimizer(learning_rate=0.0), loss_scaler)
      eval_metric_ops = metric_fn(per_example_loss, label_ids, logits)
      output_spec = tf.estimator.EstimatorSpec(
          mode=mode,
          loss=total_loss,
          eval_metric_ops=eval_metric_ops)
    else:
      dummy_op = tf.no_op()
      # Need to call mixed precision graph rewrite if fp16 to enable graph rewrite
      if FLAGS.amp:
        dummy_op = tf.train.experimental.enable_mixed_precision_graph_rewrite(
            optimization.LAMBOptimizer(learning_rate=0.0))
      output_spec = tf.estimator.EstimatorSpec(
          mode=mode, predictions=probabilities)
    return output_spec

  return model_fn


# This function is not used by this file but is still used by the Colab and
# people who depend on it.
def input_fn_builder(features, batch_size, seq_length, is_training, drop_remainder, hvd=None):
  """Creates an `input_fn` closure to be passed to Estimator."""

  all_input_ids = []
  all_input_mask = []
  all_segment_ids = []
  all_label_ids = []

  for feature in features:
    all_input_ids.append(feature.input_ids)
    all_input_mask.append(feature.input_mask)
    all_segment_ids.append(feature.segment_ids)
    all_label_ids.append(feature.label_id)

  def input_fn():
    """The actual input function."""

    num_examples = len(features)

    # This is for demo purposes and does NOT scale to large data sets. We do
    # not use Dataset.from_generator() because that uses tf.py_func which is
    # not TPU compatible. The right way to load data is with TFRecordReader.
    d = tf.data.Dataset.from_tensor_slices({
        "input_ids":
            tf.constant(
                all_input_ids, shape=[num_examples, seq_length],
                dtype=tf.int32),
        "input_mask":
            tf.constant(
                all_input_mask,
                shape=[num_examples, seq_length],
                dtype=tf.int32),
        "segment_ids":
            tf.constant(
                all_segment_ids,
                shape=[num_examples, seq_length],
                dtype=tf.int32),
        "label_ids":
            tf.constant(all_label_ids, shape=[num_examples], dtype=tf.int32),
    })

    if is_training:
      if hvd is not None: d = d.shard(hvd.size(), hvd.rank())
      d = d.repeat()
      d = d.shuffle(buffer_size=100)

    d = d.batch(batch_size=batch_size, drop_remainder=drop_remainder)
    return d

  return input_fn


def main(_):

  setup_xla_flags()

  tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.INFO)
  dllogging = utils.dllogger_class.dllogger_class(FLAGS.dllog_path)

  if FLAGS.horovod:
    hvd.init()

  processors = {
      "cola": ColaProcessor,
      "mnli": MnliProcessor,
      "mrpc": MrpcProcessor,
      "xnli": XnliProcessor,
  }

  if not FLAGS.do_train and not FLAGS.do_eval and not FLAGS.do_predict:
    raise ValueError(
        "At least one of `do_train`, `do_eval` or `do_predict' must be True.")

  bert_config = modeling.BertConfig.from_json_file(FLAGS.bert_config_file)

  if FLAGS.max_seq_length > bert_config.max_position_embeddings:
    raise ValueError(
        "Cannot use sequence length %d because the BERT model "
        "was only trained up to sequence length %d" %
        (FLAGS.max_seq_length, bert_config.max_position_embeddings))

  tf.io.gfile.makedirs(FLAGS.output_dir)

  task_name = FLAGS.task_name.lower()

  if task_name not in processors:
    raise ValueError("Task not found: %s" % (task_name))

  processor = processors[task_name]()

  label_list = processor.get_labels()

  tokenizer = tokenization.FullTokenizer(
      vocab_file=FLAGS.vocab_file, do_lower_case=FLAGS.do_lower_case)

  master_process = True
  training_hooks = []
  global_batch_size = FLAGS.train_batch_size * FLAGS.num_accumulation_steps
  hvd_rank = 0

  config = tf.compat.v1.ConfigProto()
  if FLAGS.horovod:

      tf.compat.v1.logging.info("Multi-GPU training with TF Horovod")
      tf.compat.v1.logging.info("hvd.size() = %d hvd.rank() = %d", hvd.size(), hvd.rank())
      global_batch_size = FLAGS.train_batch_size * FLAGS.num_accumulation_steps * hvd.size()
      master_process = (hvd.rank() == 0)
      hvd_rank = hvd.rank()
      config.gpu_options.visible_device_list = str(hvd.local_rank())
      set_affinity(hvd.local_rank())
      if hvd.size() > 1:
          training_hooks.append(hvd.BroadcastGlobalVariablesHook(0))
  if FLAGS.use_xla:
    config.graph_options.optimizer_options.global_jit_level = tf.compat.v1.OptimizerOptions.ON_1
    if FLAGS.amp:
      tf.enable_resource_variables()

  run_config = tf.estimator.RunConfig(
      model_dir=FLAGS.output_dir if master_process else None,
      session_config=config,
      save_checkpoints_steps=FLAGS.save_checkpoints_steps if master_process else None,
      save_summary_steps=FLAGS.save_checkpoints_steps if master_process else None,
      log_step_count_steps=FLAGS.display_loss_steps,
      keep_checkpoint_max=1)

  if master_process:
      tf.compat.v1.logging.info("***** Configuaration *****")
      for key in FLAGS.__flags.keys():
          tf.compat.v1.logging.info('  {}: {}'.format(key, getattr(FLAGS, key)))
      tf.compat.v1.logging.info("**************************")

  train_examples = None
  num_train_steps = None
  num_warmup_steps = None
  training_hooks.append(LogTrainRunHook(global_batch_size, hvd_rank, FLAGS.save_checkpoints_steps, num_steps_ignore_xla=25))

  if FLAGS.do_train:
    train_examples = processor.get_train_examples(FLAGS.data_dir)
    num_train_steps = int(
        len(train_examples) / global_batch_size * FLAGS.num_train_epochs)
    num_warmup_steps = int(num_train_steps * FLAGS.warmup_proportion)

    start_index = 0
    end_index = len(train_examples)
    tmp_filenames = [os.path.join(FLAGS.output_dir, "train.tf_record")]

    if FLAGS.horovod:
      tmp_filenames = [os.path.join(FLAGS.output_dir, "train.tf_record{}".format(i)) for i in range(hvd.size())]
      num_examples_per_rank = len(train_examples) // hvd.size()
      remainder = len(train_examples) % hvd.size()
      if hvd.rank() < remainder:
        start_index = hvd.rank() * (num_examples_per_rank+1)
        end_index = start_index + num_examples_per_rank + 1
      else:
        start_index = hvd.rank() * num_examples_per_rank + remainder
        end_index = start_index + (num_examples_per_rank)

  model_fn = model_fn_builder(
      task_name=task_name,
      bert_config=bert_config,
      num_labels=len(label_list),
      init_checkpoint=FLAGS.init_checkpoint,
      learning_rate=FLAGS.learning_rate if not FLAGS.horovod else FLAGS.learning_rate * hvd.size(),
      num_train_steps=num_train_steps,
      num_warmup_steps=num_warmup_steps,
      use_one_hot_embeddings=False,
      hvd=None if not FLAGS.horovod else hvd)

  estimator = tf.estimator.Estimator(
      model_fn=model_fn,
      config=run_config)

  if FLAGS.do_train:

    file_based_convert_examples_to_features(
        train_examples[start_index:end_index], label_list, FLAGS.max_seq_length, tokenizer, tmp_filenames[hvd_rank])

    tf.compat.v1.logging.info("***** Running training *****")
    tf.compat.v1.logging.info("  Num examples = %d", len(train_examples))
    tf.compat.v1.logging.info("  Batch size = %d", FLAGS.train_batch_size)
    tf.compat.v1.logging.info("  Num steps = %d", num_train_steps)
    train_input_fn = file_based_input_fn_builder(
        input_file=tmp_filenames,
        batch_size=FLAGS.train_batch_size,
        seq_length=FLAGS.max_seq_length,
        is_training=True,
        drop_remainder=True,
        hvd=None if not FLAGS.horovod else hvd)

    train_start_time = time.time()
    estimator.train(input_fn=train_input_fn, max_steps=num_train_steps, hooks=training_hooks)
    train_time_elapsed = time.time() - train_start_time
    train_time_wo_overhead = training_hooks[-1].total_time
    avg_sentences_per_second = num_train_steps * global_batch_size * 1.0 / train_time_elapsed
    ss_sentences_per_second = (training_hooks[-1].count - training_hooks[-1].skipped) * global_batch_size * 1.0 / train_time_wo_overhead

    if master_process:
        tf.compat.v1.logging.info("-----------------------------")
        tf.compat.v1.logging.info("Total Training Time = %0.2f for Sentences = %d", train_time_elapsed,
                        num_train_steps * global_batch_size)
        tf.compat.v1.logging.info("Total Training Time W/O Overhead = %0.2f for Sentences = %d", train_time_wo_overhead,
                        (training_hooks[-1].count - training_hooks[-1].skipped) * global_batch_size)
        tf.compat.v1.logging.info("Throughput Average (sentences/sec) with overhead = %0.2f", avg_sentences_per_second)
        tf.compat.v1.logging.info("Throughput Average (sentences/sec) = %0.2f", ss_sentences_per_second)
        tf.compat.v1.logging.info("-----------------------------")

  if FLAGS.do_eval and master_process:
    eval_examples = processor.get_dev_examples(FLAGS.data_dir)
    eval_file = os.path.join(FLAGS.output_dir, "eval.tf_record")
    file_based_convert_examples_to_features(
        eval_examples, label_list, FLAGS.max_seq_length, tokenizer, eval_file)

    tf.compat.v1.logging.info("***** Running evaluation *****")
    tf.compat.v1.logging.info("  Num examples = %d", len(eval_examples))
    tf.compat.v1.logging.info("  Batch size = %d", FLAGS.eval_batch_size)

    eval_drop_remainder = False
    eval_input_fn = file_based_input_fn_builder(
        input_file=eval_file,
        batch_size=FLAGS.eval_batch_size,
        seq_length=FLAGS.max_seq_length,
        is_training=False,
        drop_remainder=eval_drop_remainder)

    eval_hooks = [LogEvalRunHook(FLAGS.eval_batch_size)]
    eval_start_time = time.time()
    result = estimator.evaluate(input_fn=eval_input_fn, hooks=eval_hooks)

    eval_time_elapsed = time.time() - eval_start_time

    time_list = eval_hooks[-1].time_list
    time_list.sort()
    # Removing outliers (init/warmup) in throughput computation.
    eval_time_wo_overhead = sum(time_list[:int(len(time_list) * 0.8)])
    num_sentences = (int(len(time_list) * 0.8)) * FLAGS.eval_batch_size

    avg = np.mean(time_list)
    cf_50 = max(time_list[:int(len(time_list) * 0.50)])
    cf_90 = max(time_list[:int(len(time_list) * 0.90)])
    cf_95 = max(time_list[:int(len(time_list) * 0.95)])
    cf_99 = max(time_list[:int(len(time_list) * 0.99)])
    cf_100 = max(time_list[:int(len(time_list) * 1)])
    ss_sentences_per_second = num_sentences * 1.0 / eval_time_wo_overhead

    tf.compat.v1.logging.info("-----------------------------")
    tf.compat.v1.logging.info("Total Inference Time = %0.2f for Sentences = %d", eval_time_elapsed,
                    eval_hooks[-1].count * FLAGS.eval_batch_size)
    tf.compat.v1.logging.info("Total Inference Time W/O Overhead = %0.2f for Sentences = %d", eval_time_wo_overhead,
                    num_sentences)
    tf.compat.v1.logging.info("Summary Inference Statistics on EVAL set")
    tf.compat.v1.logging.info("Batch size = %d", FLAGS.eval_batch_size)
    tf.compat.v1.logging.info("Sequence Length = %d", FLAGS.max_seq_length)
    tf.compat.v1.logging.info("Precision = %s", "fp16" if FLAGS.amp else "fp32")
    tf.compat.v1.logging.info("Latency Confidence Level 50 (ms) = %0.2f", cf_50 * 1000)
    tf.compat.v1.logging.info("Latency Confidence Level 90 (ms) = %0.2f", cf_90 * 1000)
    tf.compat.v1.logging.info("Latency Confidence Level 95 (ms) = %0.2f", cf_95 * 1000)
    tf.compat.v1.logging.info("Latency Confidence Level 99 (ms) = %0.2f", cf_99 * 1000)
    tf.compat.v1.logging.info("Latency Confidence Level 100 (ms) = %0.2f", cf_100 * 1000)
    tf.compat.v1.logging.info("Latency Average (ms) = %0.2f", avg * 1000)
    tf.compat.v1.logging.info("Throughput Average (sentences/sec) = %0.2f", ss_sentences_per_second)
    dllogging.logger.log(step=(), data={"throughput_val": ss_sentences_per_second}, verbosity=Verbosity.DEFAULT)
    tf.compat.v1.logging.info("-----------------------------")


    output_eval_file = os.path.join(FLAGS.output_dir, "eval_results.txt")
    with tf.io.gfile.GFile(output_eval_file, "w") as writer:
      tf.compat.v1.logging.info("***** Eval results *****")
      for key in sorted(result.keys()):
        dllogging.logger.log(step=(), data={key: float(result[key])}, verbosity=Verbosity.DEFAULT)
        tf.compat.v1.logging.info("  %s = %s", key, str(result[key]))
        writer.write("%s = %s\n" % (key, str(result[key])))

  if FLAGS.do_predict and master_process:
    predict_examples = processor.get_test_examples(FLAGS.data_dir)
    predict_file = os.path.join(FLAGS.output_dir, "predict.tf_record")
    file_based_convert_examples_to_features(predict_examples, label_list,
                                            FLAGS.max_seq_length, tokenizer,
                                            predict_file)

    tf.compat.v1.logging.info("***** Running prediction*****")
    tf.compat.v1.logging.info("  Num examples = %d", len(predict_examples))
    tf.compat.v1.logging.info("  Batch size = %d", FLAGS.predict_batch_size)

    predict_drop_remainder = False
    predict_input_fn = file_based_input_fn_builder(
        input_file=predict_file,
        batch_size=FLAGS.predict_batch_size,
        seq_length=FLAGS.max_seq_length,
        is_training=False,
        drop_remainder=predict_drop_remainder)

    predict_hooks = [LogEvalRunHook(FLAGS.predict_batch_size)]
    predict_start_time = time.time()

    output_predict_file = os.path.join(FLAGS.output_dir, "test_results.tsv")
    with tf.io.gfile.GFile(output_predict_file, "w") as writer:
        tf.compat.v1.logging.info("***** Predict results *****")
        for prediction in estimator.predict(input_fn=predict_input_fn, hooks=predict_hooks,
                                            yield_single_examples=False):
            output_line = "\t".join(
                str(class_probability) for class_probability in prediction) + "\n"
            writer.write(output_line)


    predict_time_elapsed = time.time() - predict_start_time

    time_list = predict_hooks[-1].time_list
    time_list.sort()
    # Removing outliers (init/warmup) in throughput computation.
    predict_time_wo_overhead = sum(time_list[:int(len(time_list) * 0.8)])
    num_sentences = (int(len(time_list) * 0.8)) * FLAGS.predict_batch_size

    avg = np.mean(time_list)
    cf_50 = max(time_list[:int(len(time_list) * 0.50)])
    cf_90 = max(time_list[:int(len(time_list) * 0.90)])
    cf_95 = max(time_list[:int(len(time_list) * 0.95)])
    cf_99 = max(time_list[:int(len(time_list) * 0.99)])
    cf_100 = max(time_list[:int(len(time_list) * 1)])
    ss_sentences_per_second = num_sentences * 1.0 / predict_time_wo_overhead

    tf.compat.v1.logging.info("-----------------------------")
    tf.compat.v1.logging.info("Total Inference Time = %0.2f for Sentences = %d", predict_time_elapsed,
                    predict_hooks[-1].count * FLAGS.predict_batch_size)
    tf.compat.v1.logging.info("Total Inference Time W/O Overhead = %0.2f for Sentences = %d", predict_time_wo_overhead,
                              num_sentences)
    tf.compat.v1.logging.info("Summary Inference Statistics on TEST SET")
    tf.compat.v1.logging.info("Batch size = %d", FLAGS.predict_batch_size)
    tf.compat.v1.logging.info("Sequence Length = %d", FLAGS.max_seq_length)
    tf.compat.v1.logging.info("Precision = %s", "fp16" if FLAGS.amp else "fp32")
    tf.compat.v1.logging.info("Latency Confidence Level 50 (ms) = %0.2f", cf_50 * 1000)
    tf.compat.v1.logging.info("Latency Confidence Level 90 (ms) = %0.2f", cf_90 * 1000)
    tf.compat.v1.logging.info("Latency Confidence Level 95 (ms) = %0.2f", cf_95 * 1000)
    tf.compat.v1.logging.info("Latency Confidence Level 99 (ms) = %0.2f", cf_99 * 1000)
    tf.compat.v1.logging.info("Latency Confidence Level 100 (ms) = %0.2f", cf_100 * 1000)
    tf.compat.v1.logging.info("Latency Average (ms) = %0.2f", avg * 1000)
    tf.compat.v1.logging.info("Throughput Average (sentences/sec) = %0.2f", ss_sentences_per_second)
    dllogging.logger.log(step=(), data={"throughput_val": ss_sentences_per_second}, verbosity=Verbosity.DEFAULT)
    tf.compat.v1.logging.info("-----------------------------")


if __name__ == "__main__":
  flags.mark_flag_as_required("data_dir")
  flags.mark_flag_as_required("task_name")
  flags.mark_flag_as_required("vocab_file")
  flags.mark_flag_as_required("bert_config_file")
  flags.mark_flag_as_required("output_dir")
  tf.compat.v1.app.run()

[(184, 'modeling.BertModel', 'modeling.BertModel', False, 'import modeling\n'), (227, 'tensorflow.compat.v1.ConfigProto', 'tf.compat.v1.ConfigProto', True, 'import tensorflow as tf\n'), (460, 'utils.utils.setup_xla_flags', 'setup_xla_flags', False, 'from utils.utils import LogEvalRunHook, LogTrainRunHook, setup_xla_flags\n'), (462, 'tensorflow.compat.v1.logging.set_verbosity', 'tf.compat.v1.logging.set_verbosity', True, 'import tensorflow as tf\n'), (479, 'modeling.BertConfig.from_json_file', 'modeling.BertConfig.from_json_file', False, 'import modeling\n'), (487, 'tensorflow.io.gfile.makedirs', 'tf.io.gfile.makedirs', True, 'import tensorflow as tf\n'), (498, 'tokenization.FullTokenizer', 'tokenization.FullTokenizer', False, 'import tokenization\n'), (506, 'tensorflow.compat.v1.ConfigProto', 'tf.compat.v1.ConfigProto', True, 'import tensorflow as tf\n'), (523, 'tensorflow.estimator.RunConfig', 'tf.estimator.RunConfig', True, 'import tensorflow as tf\n'), (574, 'tensorflow.estimator.Estimator', 'tf.estimator.Estimator', True, 'import tensorflow as tf\n'), (752, 'tensorflow.compat.v1.app.run', 'tf.compat.v1.app.run', True, 'import tensorflow as tf\n'), (139, 'tensorflow.io.FixedLenFeature', 'tf.io.FixedLenFeature', True, 'import tensorflow as tf\n'), (140, 'tensorflow.io.FixedLenFeature', 'tf.io.FixedLenFeature', True, 'import tensorflow as tf\n'), (141, 'tensorflow.io.FixedLenFeature', 'tf.io.FixedLenFeature', True, 'import tensorflow as tf\n'), (142, 'tensorflow.io.FixedLenFeature', 'tf.io.FixedLenFeature', True, 'import tensorflow as tf\n'), (147, 'tensorflow.parse_single_example', 'tf.parse_single_example', True, 'import tensorflow as tf\n'), (164, 'tensorflow.data.TFRecordDataset', 'tf.data.TFRecordDataset', True, 'import tensorflow as tf\n'), (209, 'tensorflow.variable_scope', 'tf.variable_scope', True, 'import tensorflow as tf\n'), (214, 'tensorflow.matmul', 'tf.matmul', True, 'import tensorflow as tf\n'), (215, 'tensorflow.nn.bias_add', 'tf.nn.bias_add', True, 'import tensorflow as tf\n'), (216, 'tensorflow.nn.softmax', 'tf.nn.softmax', True, 'import tensorflow as tf\n'), (217, 'tensorflow.nn.log_softmax', 'tf.nn.log_softmax', True, 'import tensorflow as tf\n'), (219, 'tensorflow.one_hot', 'tf.one_hot', True, 'import tensorflow as tf\n'), (222, 'tensorflow.reduce_mean', 'tf.reduce_mean', True, 'import tensorflow as tf\n'), (231, 'tensorflow.Session', 'tf.Session', True, 'import tensorflow as tf\n'), (232, 'tensorflow.placeholder', 'tf.placeholder', True, 'import tensorflow as tf\n'), (233, 'tensorflow.placeholder', 'tf.placeholder', True, 'import tensorflow as tf\n'), (234, 'tensorflow.placeholder', 'tf.placeholder', True, 'import tensorflow as tf\n'), (235, 'tensorflow.placeholder', 'tf.placeholder', True, 'import tensorflow as tf\n'), (240, 'tensorflow.trainable_variables', 'tf.trainable_variables', True, 'import tensorflow as tf\n'), (241, 'modeling.get_assignment_map_from_checkpoint', 'modeling.get_assignment_map_from_checkpoint', False, 'import modeling\n'), (242, 'tensorflow.train.init_from_checkpoint', 'tf.train.init_from_checkpoint', True, 'import tensorflow as tf\n'), (245, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (260, 'tensorflow.python.compiler.tensorrt.trt_convert.TrtGraphConverter', 'trt.TrtGraphConverter', True, 'from tensorflow.python.compiler.tensorrt import trt_convert as trt\n'), (320, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (321, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (353, 'tensorflow.trainable_variables', 'tf.trainable_variables', True, 'import tensorflow as tf\n'), (466, 'horovod.tensorflow.init', 'hvd.init', True, 'import horovod.tensorflow as hvd\n'), (509, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (513, 'horovod.tensorflow.rank', 'hvd.rank', True, 'import horovod.tensorflow as hvd\n'), (532, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (535, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (540, 'utils.utils.LogTrainRunHook', 'LogTrainRunHook', False, 'from utils.utils import LogEvalRunHook, LogTrainRunHook, setup_xla_flags\n'), (583, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (585, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (586, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (595, 'time.time', 'time.time', False, 'import time\n'), (614, 'os.path.join', 'os.path.join', False, 'import os\n'), (618, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (620, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (631, 'time.time', 'time.time', False, 'import time\n'), (642, 'numpy.mean', 'np.mean', True, 'import numpy as np\n'), (650, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (651, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (653, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (655, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (656, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (657, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (658, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (659, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (660, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (661, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (662, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (663, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (664, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (665, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (667, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (670, 'os.path.join', 'os.path.join', False, 'import os\n'), (680, 'os.path.join', 'os.path.join', False, 'import os\n'), (685, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (687, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (698, 'time.time', 'time.time', False, 'import time\n'), (700, 'os.path.join', 'os.path.join', False, 'import os\n'), (718, 'numpy.mean', 'np.mean', True, 'import numpy as np\n'), (726, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (727, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (729, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (731, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (732, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (733, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (734, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (735, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (736, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (737, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (738, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (739, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (740, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (741, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (743, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (204, 'tensorflow.truncated_normal_initializer', 'tf.truncated_normal_initializer', True, 'import tensorflow as tf\n'), (207, 'tensorflow.zeros_initializer', 'tf.zeros_initializer', True, 'import tensorflow as tf\n'), (212, 'tensorflow.nn.dropout', 'tf.nn.dropout', True, 'import tensorflow as tf\n'), (221, 'tensorflow.reduce_sum', 'tf.reduce_sum', True, 'import tensorflow as tf\n'), (243, 'tensorflow.global_variables_initializer', 'tf.global_variables_initializer', True, 'import tensorflow as tf\n'), (276, 'tensorflow.io.gfile.GFile', 'tf.io.gfile.GFile', True, 'import tensorflow as tf\n'), (292, 'tensorflow.argmax', 'tf.argmax', True, 'import tensorflow as tf\n'), (323, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (334, 'tensorflow.import_graph_def', 'tf.import_graph_def', True, 'import tensorflow as tf\n'), (357, 'modeling.get_assignment_map_from_checkpoint', 'modeling.get_assignment_map_from_checkpoint', False, 'import modeling\n'), (358, 'tensorflow.train.init_from_checkpoint', 'tf.train.init_from_checkpoint', True, 'import tensorflow as tf\n'), (361, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (372, 'optimization.create_optimizer', 'optimization.create_optimizer', False, 'import optimization\n'), (375, 'tensorflow.estimator.EstimatorSpec', 'tf.estimator.EstimatorSpec', True, 'import tensorflow as tf\n'), (510, 'horovod.tensorflow.size', 'hvd.size', True, 'import horovod.tensorflow as hvd\n'), (510, 'horovod.tensorflow.rank', 'hvd.rank', True, 'import horovod.tensorflow as hvd\n'), (511, 'horovod.tensorflow.size', 'hvd.size', True, 'import horovod.tensorflow as hvd\n'), (512, 'horovod.tensorflow.rank', 'hvd.rank', True, 'import horovod.tensorflow as hvd\n'), (514, 'horovod.tensorflow.local_rank', 'hvd.local_rank', True, 'import horovod.tensorflow as hvd\n'), (515, 'horovod.tensorflow.local_rank', 'hvd.local_rank', True, 'import horovod.tensorflow as hvd\n'), (516, 'horovod.tensorflow.size', 'hvd.size', True, 'import horovod.tensorflow as hvd\n'), (521, 'tensorflow.enable_resource_variables', 'tf.enable_resource_variables', True, 'import tensorflow as tf\n'), (550, 'os.path.join', 'os.path.join', False, 'import os\n'), (597, 'time.time', 'time.time', False, 'import time\n'), (603, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (604, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (606, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (608, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (609, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (610, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (630, 'utils.utils.LogEvalRunHook', 'LogEvalRunHook', False, 'from utils.utils import LogEvalRunHook, LogTrainRunHook, setup_xla_flags\n'), (634, 'time.time', 'time.time', False, 'import time\n'), (671, 'tensorflow.io.gfile.GFile', 'tf.io.gfile.GFile', True, 'import tensorflow as tf\n'), (672, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (697, 'utils.utils.LogEvalRunHook', 'LogEvalRunHook', False, 'from utils.utils import LogEvalRunHook, LogTrainRunHook, setup_xla_flags\n'), (701, 'tensorflow.io.gfile.GFile', 'tf.io.gfile.GFile', True, 'import tensorflow as tf\n'), (702, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (710, 'time.time', 'time.time', False, 'import time\n'), (154, 'tensorflow.to_int32', 'tf.to_int32', True, 'import tensorflow as tf\n'), (252, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (294, 'tensorflow.metrics.false_negatives', 'tf.metrics.false_negatives', True, 'import tensorflow as tf\n'), (295, 'tensorflow.metrics.false_positives', 'tf.metrics.false_positives', True, 'import tensorflow as tf\n'), (296, 'tensorflow.metrics.true_positives', 'tf.metrics.true_positives', True, 'import tensorflow as tf\n'), (297, 'tensorflow.metrics.true_negatives', 'tf.metrics.true_negatives', True, 'import tensorflow as tf\n'), (340, 'tensorflow.estimator.EstimatorSpec', 'tf.estimator.EstimatorSpec', True, 'import tensorflow as tf\n'), (366, 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', True, 'import tensorflow as tf\n'), (380, 'tensorflow.no_op', 'tf.no_op', True, 'import tensorflow as tf\n'), (387, 'tensorflow.estimator.EstimatorSpec', 'tf.estimator.EstimatorSpec', True, 'import tensorflow as tf\n'), (392, 'tensorflow.no_op', 'tf.no_op', True, 'import tensorflow as tf\n'), (397, 'tensorflow.estimator.EstimatorSpec', 'tf.estimator.EstimatorSpec', True, 'import tensorflow as tf\n'), (430, 'tensorflow.constant', 'tf.constant', True, 'import tensorflow as tf\n'), (434, 'tensorflow.constant', 'tf.constant', True, 'import tensorflow as tf\n'), (439, 'tensorflow.constant', 'tf.constant', True, 'import tensorflow as tf\n'), (444, 'tensorflow.constant', 'tf.constant', True, 'import tensorflow as tf\n'), (517, 'horovod.tensorflow.BroadcastGlobalVariablesHook', 'hvd.BroadcastGlobalVariablesHook', True, 'import horovod.tensorflow as hvd\n'), (554, 'horovod.tensorflow.size', 'hvd.size', True, 'import horovod.tensorflow as hvd\n'), (555, 'horovod.tensorflow.size', 'hvd.size', True, 'import horovod.tensorflow as hvd\n'), (556, 'horovod.tensorflow.rank', 'hvd.rank', True, 'import horovod.tensorflow as hvd\n'), (166, 'horovod.tensorflow.size', 'hvd.size', True, 'import horovod.tensorflow as hvd\n'), (166, 'horovod.tensorflow.rank', 'hvd.rank', True, 'import horovod.tensorflow as hvd\n'), (300, 'tensorflow.identity', 'tf.identity', True, 'import tensorflow as tf\n'), (303, 'tensorflow.metrics.accuracy', 'tf.metrics.accuracy', True, 'import tensorflow as tf\n'), (305, 'tensorflow.metrics.mean', 'tf.metrics.mean', True, 'import tensorflow as tf\n'), (306, 'tf_metrics.f1', 'tf_metrics.f1', False, 'import tf_metrics\n'), (313, 'tensorflow.metrics.accuracy', 'tf.metrics.accuracy', True, 'import tensorflow as tf\n'), (315, 'tensorflow.metrics.mean', 'tf.metrics.mean', True, 'import tensorflow as tf\n'), (344, 'tensorflow.estimator.EstimatorSpec', 'tf.estimator.EstimatorSpec', True, 'import tensorflow as tf\n'), (355, 'horovod.tensorflow.rank', 'hvd.rank', True, 'import horovod.tensorflow as hvd\n'), (383, 'tensorflow.train.experimental.FixedLossScale', 'tf.train.experimental.FixedLossScale', True, 'import tensorflow as tf\n'), (448, 'horovod.tensorflow.size', 'hvd.size', True, 'import horovod.tensorflow as hvd\n'), (448, 'horovod.tensorflow.rank', 'hvd.rank', True, 'import horovod.tensorflow as hvd\n'), (557, 'horovod.tensorflow.rank', 'hvd.rank', True, 'import horovod.tensorflow as hvd\n'), (568, 'horovod.tensorflow.size', 'hvd.size', True, 'import horovod.tensorflow as hvd\n'), (385, 'optimization.LAMBOptimizer', 'optimization.LAMBOptimizer', False, 'import optimization\n'), (396, 'optimization.LAMBOptimizer', 'optimization.LAMBOptimizer', False, 'import optimization\n'), (553, 'horovod.tensorflow.size', 'hvd.size', True, 'import horovod.tensorflow as hvd\n'), (560, 'horovod.tensorflow.rank', 'hvd.rank', True, 'import horovod.tensorflow as hvd\n')]