# 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()

[((6344, 6560), 'modeling.BertModel', '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'}), '(config=bert_config, is_training=is_training, input_ids=\n    input_ids, input_mask=input_mask, token_type_ids=segment_ids,\n    use_one_hot_embeddings=use_one_hot_embeddings, compute_type=tf.float32)\n', (6362, 6560), False, 'import modeling\n'), ((7956, 7982), 'tensorflow.compat.v1.ConfigProto', 'tf.compat.v1.ConfigProto', ([], {}), '()\n', (7980, 7982), True, 'import tensorflow as tf\n'), ((17775, 17792), 'utils.utils.setup_xla_flags', 'setup_xla_flags', ([], {}), '()\n', (17790, 17792), False, 'from utils.utils import LogEvalRunHook, LogTrainRunHook, setup_xla_flags\n'), ((17796, 17857), 'tensorflow.compat.v1.logging.set_verbosity', 'tf.compat.v1.logging.set_verbosity', (['tf.compat.v1.logging.INFO'], {}), '(tf.compat.v1.logging.INFO)\n', (17830, 17857), True, 'import tensorflow as tf\n'), ((18291, 18349), 'modeling.BertConfig.from_json_file', 'modeling.BertConfig.from_json_file', (['FLAGS.bert_config_file'], {}), '(FLAGS.bert_config_file)\n', (18325, 18349), False, 'import modeling\n'), ((18628, 18666), 'tensorflow.io.gfile.makedirs', 'tf.io.gfile.makedirs', (['FLAGS.output_dir'], {}), '(FLAGS.output_dir)\n', (18648, 18666), True, 'import tensorflow as tf\n'), ((18891, 18986), 'tokenization.FullTokenizer', 'tokenization.FullTokenizer', ([], {'vocab_file': 'FLAGS.vocab_file', 'do_lower_case': 'FLAGS.do_lower_case'}), '(vocab_file=FLAGS.vocab_file, do_lower_case=FLAGS\n    .do_lower_case)\n', (18917, 18986), False, 'import tokenization\n'), ((19139, 19165), 'tensorflow.compat.v1.ConfigProto', 'tf.compat.v1.ConfigProto', ([], {}), '()\n', (19163, 19165), True, 'import tensorflow as tf\n'), ((19898, 20245), 'tensorflow.estimator.RunConfig', '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)'}), '(model_dir=FLAGS.output_dir if master_process else\n    None, session_config=config, save_checkpoints_steps=FLAGS.\n    save_checkpoints_steps if master_process else None, save_summary_steps=\n    FLAGS.save_checkpoints_steps if master_process else None,\n    log_step_count_steps=FLAGS.display_loss_steps, keep_checkpoint_max=1)\n', (19920, 20245), True, 'import tensorflow as tf\n'), ((22110, 22170), 'tensorflow.estimator.Estimator', 'tf.estimator.Estimator', ([], {'model_fn': 'model_fn', 'config': 'run_config'}), '(model_fn=model_fn, config=run_config)\n', (22132, 22170), True, 'import tensorflow as tf\n'), ((31757, 31779), 'tensorflow.compat.v1.app.run', 'tf.compat.v1.app.run', ([], {}), '()\n', (31777, 31779), True, 'import tensorflow as tf\n'), ((4847, 4892), 'tensorflow.io.FixedLenFeature', 'tf.io.FixedLenFeature', (['[seq_length]', 'tf.int64'], {}), '([seq_length], tf.int64)\n', (4868, 4892), True, 'import tensorflow as tf\n'), ((4914, 4959), 'tensorflow.io.FixedLenFeature', 'tf.io.FixedLenFeature', (['[seq_length]', 'tf.int64'], {}), '([seq_length], tf.int64)\n', (4935, 4959), True, 'import tensorflow as tf\n'), ((4982, 5027), 'tensorflow.io.FixedLenFeature', 'tf.io.FixedLenFeature', (['[seq_length]', 'tf.int64'], {}), '([seq_length], tf.int64)\n', (5003, 5027), True, 'import tensorflow as tf\n'), ((5048, 5083), 'tensorflow.io.FixedLenFeature', 'tf.io.FixedLenFeature', (['[]', 'tf.int64'], {}), '([], tf.int64)\n', (5069, 5083), True, 'import tensorflow as tf\n'), ((5204, 5253), 'tensorflow.parse_single_example', 'tf.parse_single_example', (['record', 'name_to_features'], {}), '(record, name_to_features)\n', (5227, 5253), True, 'import tensorflow as tf\n'), ((5737, 5772), 'tensorflow.data.TFRecordDataset', 'tf.data.TFRecordDataset', (['input_file'], {}), '(input_file)\n', (5760, 5772), True, 'import tensorflow as tf\n'), ((7132, 7157), 'tensorflow.variable_scope', 'tf.variable_scope', (['"""loss"""'], {}), "('loss')\n", (7149, 7157), True, 'import tensorflow as tf\n'), ((7283, 7340), 'tensorflow.matmul', 'tf.matmul', (['output_layer', 'output_weights'], {'transpose_b': '(True)'}), '(output_layer, output_weights, transpose_b=True)\n', (7292, 7340), True, 'import tensorflow as tf\n'), ((7354, 7408), 'tensorflow.nn.bias_add', 'tf.nn.bias_add', (['logits', 'output_bias'], {'name': '"""cls_logits"""'}), "(logits, output_bias, name='cls_logits')\n", (7368, 7408), True, 'import tensorflow as tf\n'), ((7429, 7485), 'tensorflow.nn.softmax', 'tf.nn.softmax', (['logits'], {'axis': '(-1)', 'name': '"""cls_probabilities"""'}), "(logits, axis=-1, name='cls_probabilities')\n", (7442, 7485), True, 'import tensorflow as tf\n'), ((7502, 7536), 'tensorflow.nn.log_softmax', 'tf.nn.log_softmax', (['logits'], {'axis': '(-1)'}), '(logits, axis=-1)\n', (7519, 7536), True, 'import tensorflow as tf\n'), ((7559, 7613), 'tensorflow.one_hot', 'tf.one_hot', (['labels'], {'depth': 'num_labels', 'dtype': 'tf.float32'}), '(labels, depth=num_labels, dtype=tf.float32)\n', (7569, 7613), True, 'import tensorflow as tf\n'), ((7730, 7779), 'tensorflow.reduce_mean', 'tf.reduce_mean', (['per_example_loss'], {'name': '"""cls_loss"""'}), "(per_example_loss, name='cls_loss')\n", (7744, 7779), True, 'import tensorflow as tf\n'), ((8149, 8177), 'tensorflow.Session', 'tf.Session', ([], {'config': 'tf_config'}), '(config=tf_config)\n', (8159, 8177), True, 'import tensorflow as tf\n'), ((8206, 8250), 'tensorflow.placeholder', 'tf.placeholder', (['tf.int32', 'shape', '"""input_ids"""'], {}), "(tf.int32, shape, 'input_ids')\n", (8220, 8250), True, 'import tensorflow as tf\n'), ((8268, 8313), 'tensorflow.placeholder', 'tf.placeholder', (['tf.int32', 'shape', '"""input_mask"""'], {}), "(tf.int32, shape, 'input_mask')\n", (8282, 8313), True, 'import tensorflow as tf\n'), ((8332, 8378), 'tensorflow.placeholder', 'tf.placeholder', (['tf.int32', 'shape', '"""segment_ids"""'], {}), "(tf.int32, shape, 'segment_ids')\n", (8346, 8378), True, 'import tensorflow as tf\n'), ((8395, 8438), 'tensorflow.placeholder', 'tf.placeholder', (['tf.int32', 'None', '"""label_ids"""'], {}), "(tf.int32, None, 'label_ids')\n", (8409, 8438), True, 'import tensorflow as tf\n'), ((8587, 8611), 'tensorflow.trainable_variables', 'tf.trainable_variables', ([], {}), '()\n', (8609, 8611), True, 'import tensorflow as tf\n'), ((8663, 8730), 'modeling.get_assignment_map_from_checkpoint', 'modeling.get_assignment_map_from_checkpoint', (['tvars', 'init_checkpoint'], {}), '(tvars, init_checkpoint)\n', (8706, 8730), False, 'import modeling\n'), ((8735, 8797), 'tensorflow.train.init_from_checkpoint', 'tf.train.init_from_checkpoint', (['init_checkpoint', 'assignment_map'], {}), '(init_checkpoint, assignment_map)\n', (8764, 8797), True, 'import tensorflow as tf\n'), ((8874, 8932), 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', (['"""**** Trainable Variables ****"""'], {}), "('**** Trainable Variables ****')\n", (8899, 8932), True, 'import tensorflow as tf\n'), ((9552, 9818), 'tensorflow.python.compiler.tensorrt.trt_convert.TrtGraphConverter', '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)'}), "(input_graph_def=frozen_graph, nodes_blacklist=\n    output_node_names, max_workspace_size_bytes=(4096 << 20) - 1000,\n    precision_mode='FP16' if FLAGS.amp else 'FP32', minimum_segment_size=4,\n    is_dynamic_op=True, maximum_cached_engines=1000)\n", (9573, 9818), True, 'from tensorflow.python.compiler.tensorrt import trt_convert as trt\n'), ((12163, 12208), 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', (['"""*** Features ***"""'], {}), "('*** Features ***')\n", (12188, 12208), True, 'import tensorflow as tf\n'), ((12213, 12258), 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', (['"""*** Features ***"""'], {}), "('*** Features ***')\n", (12238, 12258), True, 'import tensorflow as tf\n'), ((13898, 13922), 'tensorflow.trainable_variables', 'tf.trainable_variables', ([], {}), '()\n', (13920, 13922), True, 'import tensorflow as tf\n'), ((17951, 17961), 'horovod.tensorflow.init', 'hvd.init', ([], {}), '()\n', (17959, 17961), True, 'import horovod.tensorflow as hvd\n'), ((19193, 19256), 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', (['"""Multi-GPU training with TF Horovod"""'], {}), "('Multi-GPU training with TF Horovod')\n", (19218, 19256), True, 'import tensorflow as tf\n'), ((19499, 19509), 'horovod.tensorflow.rank', 'hvd.rank', ([], {}), '()\n', (19507, 19509), True, 'import horovod.tensorflow as hvd\n'), ((20293, 20348), 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', (['"""***** Configuaration *****"""'], {}), "('***** Configuaration *****')\n", (20318, 20348), True, 'import tensorflow as tf\n'), ((20475, 20530), 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', (['"""**************************"""'], {}), "('**************************')\n", (20500, 20530), True, 'import tensorflow as tf\n'), ((20631, 20734), 'utils.utils.LogTrainRunHook', 'LogTrainRunHook', (['global_batch_size', 'hvd_rank', 'FLAGS.save_checkpoints_steps'], {'num_steps_ignore_xla': '(25)'}), '(global_batch_size, hvd_rank, FLAGS.save_checkpoints_steps,\n    num_steps_ignore_xla=25)\n', (20646, 20734), False, 'from utils.utils import LogEvalRunHook, LogTrainRunHook, setup_xla_flags\n'), ((22374, 22431), 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', (['"""***** Running training *****"""'], {}), "('***** Running training *****')\n", (22399, 22431), True, 'import tensorflow as tf\n'), ((22510, 22580), 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', (['"""  Batch size = %d"""', 'FLAGS.train_batch_size'], {}), "('  Batch size = %d', FLAGS.train_batch_size)\n", (22535, 22580), True, 'import tensorflow as tf\n'), ((22585, 22647), 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', (['"""  Num steps = %d"""', 'num_train_steps'], {}), "('  Num steps = %d', num_train_steps)\n", (22610, 22647), True, 'import tensorflow as tf\n'), ((22943, 22954), 'time.time', 'time.time', ([], {}), '()\n', (22952, 22954), False, 'import time\n'), ((24284, 24332), 'os.path.join', 'os.path.join', (['FLAGS.output_dir', '"""eval.tf_record"""'], {}), "(FLAGS.output_dir, 'eval.tf_record')\n", (24296, 24332), False, 'import os\n'), ((24462, 24521), 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', (['"""***** Running evaluation *****"""'], {}), "('***** Running evaluation *****')\n", (24487, 24521), True, 'import tensorflow as tf\n'), ((24599, 24668), 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', (['"""  Batch size = %d"""', 'FLAGS.eval_batch_size'], {}), "('  Batch size = %d', FLAGS.eval_batch_size)\n", (24624, 24668), True, 'import tensorflow as tf\n'), ((25015, 25026), 'time.time', 'time.time', ([], {}), '()\n', (25024, 25026), False, 'import time\n'), ((25438, 25456), 'numpy.mean', 'np.mean', (['time_list'], {}), '(time_list)\n', (25445, 25456), True, 'import numpy as np\n'), ((25814, 25872), 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', (['"""-----------------------------"""'], {}), "('-----------------------------')\n", (25839, 25872), True, 'import tensorflow as tf\n'), ((25877, 26022), 'tensorflow.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)'], {}), "('Total Inference Time = %0.2f for Sentences = %d',\n    eval_time_elapsed, eval_hooks[-1].count * FLAGS.eval_batch_size)\n", (25902, 26022), True, 'import tensorflow as tf\n'), ((26043, 26179), 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', (['"""Total Inference Time W/O Overhead = %0.2f for Sentences = %d"""', 'eval_time_wo_overhead', 'num_sentences'], {}), "(\n    'Total Inference Time W/O Overhead = %0.2f for Sentences = %d',\n    eval_time_wo_overhead, num_sentences)\n", (26068, 26179), True, 'import tensorflow as tf\n'), ((26195, 26264), 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', (['"""Summary Inference Statistics on EVAL set"""'], {}), "('Summary Inference Statistics on EVAL set')\n", (26220, 26264), True, 'import tensorflow as tf\n'), ((26269, 26336), 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', (['"""Batch size = %d"""', 'FLAGS.eval_batch_size'], {}), "('Batch size = %d', FLAGS.eval_batch_size)\n", (26294, 26336), True, 'import tensorflow as tf\n'), ((26341, 26412), 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', (['"""Sequence Length = %d"""', 'FLAGS.max_seq_length'], {}), "('Sequence Length = %d', FLAGS.max_seq_length)\n", (26366, 26412), True, 'import tensorflow as tf\n'), ((26417, 26493), 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', (['"""Precision = %s"""', "('fp16' if FLAGS.amp else 'fp32')"], {}), "('Precision = %s', 'fp16' if FLAGS.amp else 'fp32')\n", (26442, 26493), True, 'import tensorflow as tf\n'), ((26498, 26585), 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', (['"""Latency Confidence Level 50 (ms) = %0.2f"""', '(cf_50 * 1000)'], {}), "('Latency Confidence Level 50 (ms) = %0.2f', cf_50 *\n    1000)\n", (26523, 26585), True, 'import tensorflow as tf\n'), ((26586, 26673), 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', (['"""Latency Confidence Level 90 (ms) = %0.2f"""', '(cf_90 * 1000)'], {}), "('Latency Confidence Level 90 (ms) = %0.2f', cf_90 *\n    1000)\n", (26611, 26673), True, 'import tensorflow as tf\n'), ((26674, 26761), 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', (['"""Latency Confidence Level 95 (ms) = %0.2f"""', '(cf_95 * 1000)'], {}), "('Latency Confidence Level 95 (ms) = %0.2f', cf_95 *\n    1000)\n", (26699, 26761), True, 'import tensorflow as tf\n'), ((26762, 26849), 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', (['"""Latency Confidence Level 99 (ms) = %0.2f"""', '(cf_99 * 1000)'], {}), "('Latency Confidence Level 99 (ms) = %0.2f', cf_99 *\n    1000)\n", (26787, 26849), True, 'import tensorflow as tf\n'), ((26850, 26940), 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', (['"""Latency Confidence Level 100 (ms) = %0.2f"""', '(cf_100 * 1000)'], {}), "('Latency Confidence Level 100 (ms) = %0.2f', \n    cf_100 * 1000)\n", (26875, 26940), True, 'import tensorflow as tf\n'), ((26940, 27009), 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', (['"""Latency Average (ms) = %0.2f"""', '(avg * 1000)'], {}), "('Latency Average (ms) = %0.2f', avg * 1000)\n", (26965, 27009), True, 'import tensorflow as tf\n'), ((27014, 27114), 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', (['"""Throughput Average (sentences/sec) = %0.2f"""', 'ss_sentences_per_second'], {}), "('Throughput Average (sentences/sec) = %0.2f',\n    ss_sentences_per_second)\n", (27039, 27114), True, 'import tensorflow as tf\n'), ((27228, 27286), 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', (['"""-----------------------------"""'], {}), "('-----------------------------')\n", (27253, 27286), True, 'import tensorflow as tf\n'), ((27312, 27362), 'os.path.join', 'os.path.join', (['FLAGS.output_dir', '"""eval_results.txt"""'], {}), "(FLAGS.output_dir, 'eval_results.txt')\n", (27324, 27362), False, 'import os\n'), ((27882, 27933), 'os.path.join', 'os.path.join', (['FLAGS.output_dir', '"""predict.tf_record"""'], {}), "(FLAGS.output_dir, 'predict.tf_record')\n", (27894, 27933), False, 'import os\n'), ((28148, 28206), 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', (['"""***** Running prediction*****"""'], {}), "('***** Running prediction*****')\n", (28173, 28206), True, 'import tensorflow as tf\n'), ((28287, 28359), 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', (['"""  Batch size = %d"""', 'FLAGS.predict_batch_size'], {}), "('  Batch size = %d', FLAGS.predict_batch_size)\n", (28312, 28359), True, 'import tensorflow as tf\n'), ((28730, 28741), 'time.time', 'time.time', ([], {}), '()\n', (28739, 28741), False, 'import time\n'), ((28769, 28819), 'os.path.join', 'os.path.join', (['FLAGS.output_dir', '"""test_results.tsv"""'], {}), "(FLAGS.output_dir, 'test_results.tsv')\n", (28781, 28819), False, 'import os\n'), ((29626, 29644), 'numpy.mean', 'np.mean', (['time_list'], {}), '(time_list)\n', (29633, 29644), True, 'import numpy as np\n'), ((30005, 30063), 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', (['"""-----------------------------"""'], {}), "('-----------------------------')\n", (30030, 30063), True, 'import tensorflow as tf\n'), ((30068, 30222), 'tensorflow.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)'], {}), "('Total Inference Time = %0.2f for Sentences = %d',\n    predict_time_elapsed, predict_hooks[-1].count * FLAGS.predict_batch_size)\n", (30093, 30222), True, 'import tensorflow as tf\n'), ((30243, 30382), 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', (['"""Total Inference Time W/O Overhead = %0.2f for Sentences = %d"""', 'predict_time_wo_overhead', 'num_sentences'], {}), "(\n    'Total Inference Time W/O Overhead = %0.2f for Sentences = %d',\n    predict_time_wo_overhead, num_sentences)\n", (30268, 30382), True, 'import tensorflow as tf\n'), ((30408, 30477), 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', (['"""Summary Inference Statistics on TEST SET"""'], {}), "('Summary Inference Statistics on TEST SET')\n", (30433, 30477), True, 'import tensorflow as tf\n'), ((30482, 30552), 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', (['"""Batch size = %d"""', 'FLAGS.predict_batch_size'], {}), "('Batch size = %d', FLAGS.predict_batch_size)\n", (30507, 30552), True, 'import tensorflow as tf\n'), ((30557, 30628), 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', (['"""Sequence Length = %d"""', 'FLAGS.max_seq_length'], {}), "('Sequence Length = %d', FLAGS.max_seq_length)\n", (30582, 30628), True, 'import tensorflow as tf\n'), ((30633, 30709), 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', (['"""Precision = %s"""', "('fp16' if FLAGS.amp else 'fp32')"], {}), "('Precision = %s', 'fp16' if FLAGS.amp else 'fp32')\n", (30658, 30709), True, 'import tensorflow as tf\n'), ((30714, 30801), 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', (['"""Latency Confidence Level 50 (ms) = %0.2f"""', '(cf_50 * 1000)'], {}), "('Latency Confidence Level 50 (ms) = %0.2f', cf_50 *\n    1000)\n", (30739, 30801), True, 'import tensorflow as tf\n'), ((30802, 30889), 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', (['"""Latency Confidence Level 90 (ms) = %0.2f"""', '(cf_90 * 1000)'], {}), "('Latency Confidence Level 90 (ms) = %0.2f', cf_90 *\n    1000)\n", (30827, 30889), True, 'import tensorflow as tf\n'), ((30890, 30977), 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', (['"""Latency Confidence Level 95 (ms) = %0.2f"""', '(cf_95 * 1000)'], {}), "('Latency Confidence Level 95 (ms) = %0.2f', cf_95 *\n    1000)\n", (30915, 30977), True, 'import tensorflow as tf\n'), ((30978, 31065), 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', (['"""Latency Confidence Level 99 (ms) = %0.2f"""', '(cf_99 * 1000)'], {}), "('Latency Confidence Level 99 (ms) = %0.2f', cf_99 *\n    1000)\n", (31003, 31065), True, 'import tensorflow as tf\n'), ((31066, 31156), 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', (['"""Latency Confidence Level 100 (ms) = %0.2f"""', '(cf_100 * 1000)'], {}), "('Latency Confidence Level 100 (ms) = %0.2f', \n    cf_100 * 1000)\n", (31091, 31156), True, 'import tensorflow as tf\n'), ((31156, 31225), 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', (['"""Latency Average (ms) = %0.2f"""', '(avg * 1000)'], {}), "('Latency Average (ms) = %0.2f', avg * 1000)\n", (31181, 31225), True, 'import tensorflow as tf\n'), ((31230, 31330), 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', (['"""Throughput Average (sentences/sec) = %0.2f"""', 'ss_sentences_per_second'], {}), "('Throughput Average (sentences/sec) = %0.2f',\n    ss_sentences_per_second)\n", (31255, 31330), True, 'import tensorflow as tf\n'), ((31444, 31502), 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', (['"""-----------------------------"""'], {}), "('-----------------------------')\n", (31469, 31502), True, 'import tensorflow as tf\n'), ((6973, 7017), 'tensorflow.truncated_normal_initializer', 'tf.truncated_normal_initializer', ([], {'stddev': '(0.02)'}), '(stddev=0.02)\n', (7004, 7017), True, 'import tensorflow as tf\n'), ((7100, 7122), 'tensorflow.zeros_initializer', 'tf.zeros_initializer', ([], {}), '()\n', (7120, 7122), True, 'import tensorflow as tf\n'), ((7226, 7268), 'tensorflow.nn.dropout', 'tf.nn.dropout', (['output_layer'], {'keep_prob': '(0.9)'}), '(output_layer, keep_prob=0.9)\n', (7239, 7268), True, 'import tensorflow as tf\n'), ((7639, 7718), 'tensorflow.reduce_sum', 'tf.reduce_sum', (['(one_hot_labels * log_probs)'], {'axis': '(-1)', 'name': '"""cls_per_example_loss"""'}), "(one_hot_labels * log_probs, axis=-1, name='cls_per_example_loss')\n", (7652, 7718), True, 'import tensorflow as tf\n'), ((8814, 8847), 'tensorflow.global_variables_initializer', 'tf.global_variables_initializer', ([], {}), '()\n', (8845, 8847), True, 'import tensorflow as tf\n'), ((10147, 10192), 'tensorflow.io.gfile.GFile', 'tf.io.gfile.GFile', (['"""frozen_modelTRT.pb"""', '"""wb"""'], {}), "('frozen_modelTRT.pb', 'wb')\n", (10164, 10192), True, 'import tensorflow as tf\n'), ((10741, 10789), 'tensorflow.argmax', 'tf.argmax', (['logits'], {'axis': '(-1)', 'output_type': 'tf.int32'}), '(logits, axis=-1, output_type=tf.int32)\n', (10750, 10789), True, 'import tensorflow as tf\n'), ((12306, 12394), 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', (["('  name = %s, shape = %s' % (name, features[name].shape))"], {}), "('  name = %s, shape = %s' % (name, features[name]\n    .shape))\n", (12331, 12394), True, 'import tensorflow as tf\n'), ((12840, 13126), 'tensorflow.import_graph_def', 'tf.import_graph_def', (['trt_graph'], {'input_map': "{'input_ids': input_ids, 'input_mask': input_mask, 'segment_ids':\n    segment_ids, 'label_ids': label_ids}", 'return_elements': "['loss/cls_loss:0', 'loss/cls_per_example_loss:0', 'loss/cls_logits:0',\n    'loss/cls_probabilities:0']", 'name': '""""""'}), "(trt_graph, input_map={'input_ids': input_ids,\n    'input_mask': input_mask, 'segment_ids': segment_ids, 'label_ids':\n    label_ids}, return_elements=['loss/cls_loss:0',\n    'loss/cls_per_example_loss:0', 'loss/cls_logits:0',\n    'loss/cls_probabilities:0'], name='')\n", (12859, 13126), True, 'import tensorflow as tf\n'), ((14080, 14147), 'modeling.get_assignment_map_from_checkpoint', 'modeling.get_assignment_map_from_checkpoint', (['tvars', 'init_checkpoint'], {}), '(tvars, init_checkpoint)\n', (14123, 14147), False, 'import modeling\n'), ((14154, 14216), 'tensorflow.train.init_from_checkpoint', 'tf.train.init_from_checkpoint', (['init_checkpoint', 'assignment_map'], {}), '(init_checkpoint, assignment_map)\n', (14183, 14216), True, 'import tensorflow as tf\n'), ((14256, 14314), 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', (['"""**** Trainable Variables ****"""'], {}), "('**** Trainable Variables ****')\n", (14281, 14314), True, 'import tensorflow as tf\n'), ((14679, 14853), 'optimization.create_optimizer', 'optimization.create_optimizer', (['total_loss', 'learning_rate', 'num_train_steps', 'num_warmup_steps', 'hvd', '(False)', 'FLAGS.amp', 'FLAGS.num_accumulation_steps', 'FLAGS.optimizer_type'], {}), '(total_loss, learning_rate, num_train_steps,\n    num_warmup_steps, hvd, False, FLAGS.amp, FLAGS.num_accumulation_steps,\n    FLAGS.optimizer_type)\n', (14708, 14853), False, 'import optimization\n'), ((14887, 14960), 'tensorflow.estimator.EstimatorSpec', 'tf.estimator.EstimatorSpec', ([], {'mode': 'mode', 'loss': 'total_loss', 'train_op': 'train_op'}), '(mode=mode, loss=total_loss, train_op=train_op)\n', (14913, 14960), True, 'import tensorflow as tf\n'), ((19324, 19334), 'horovod.tensorflow.size', 'hvd.size', ([], {}), '()\n', (19332, 19334), True, 'import horovod.tensorflow as hvd\n'), ((19336, 19346), 'horovod.tensorflow.rank', 'hvd.rank', ([], {}), '()\n', (19344, 19346), True, 'import horovod.tensorflow as hvd\n'), ((19430, 19440), 'horovod.tensorflow.size', 'hvd.size', ([], {}), '()\n', (19438, 19440), True, 'import horovod.tensorflow as hvd\n'), ((19465, 19475), 'horovod.tensorflow.rank', 'hvd.rank', ([], {}), '()\n', (19473, 19475), True, 'import horovod.tensorflow as hvd\n'), ((19561, 19577), 'horovod.tensorflow.local_rank', 'hvd.local_rank', ([], {}), '()\n', (19575, 19577), True, 'import horovod.tensorflow as hvd\n'), ((19598, 19614), 'horovod.tensorflow.local_rank', 'hvd.local_rank', ([], {}), '()\n', (19612, 19614), True, 'import horovod.tensorflow as hvd\n'), ((19625, 19635), 'horovod.tensorflow.size', 'hvd.size', ([], {}), '()\n', (19633, 19635), True, 'import horovod.tensorflow as hvd\n'), ((19851, 19881), 'tensorflow.enable_resource_variables', 'tf.enable_resource_variables', ([], {}), '()\n', (19879, 19881), True, 'import tensorflow as tf\n'), ((21069, 21118), 'os.path.join', 'os.path.join', (['FLAGS.output_dir', '"""train.tf_record"""'], {}), "(FLAGS.output_dir, 'train.tf_record')\n", (21081, 21118), False, 'import os\n'), ((23074, 23085), 'time.time', 'time.time', ([], {}), '()\n', (23083, 23085), False, 'import time\n'), ((23427, 23485), 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', (['"""-----------------------------"""'], {}), "('-----------------------------')\n", (23452, 23485), True, 'import tensorflow as tf\n'), ((23494, 23630), 'tensorflow.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)'], {}), "('Total Training Time = %0.2f for Sentences = %d',\n    train_time_elapsed, num_train_steps * global_batch_size)\n", (23519, 23630), True, 'import tensorflow as tf\n'), ((23659, 23862), 'tensorflow.compat.v1.logging.info', '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)'], {}), "(\n    'Total Training Time W/O Overhead = %0.2f for Sentences = %d',\n    train_time_wo_overhead, (training_hooks[-1].count - training_hooks[-1].\n    skipped) * global_batch_size)\n", (23684, 23862), True, 'import tensorflow as tf\n'), ((23881, 24001), 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', (['"""Throughput Average (sentences/sec) with overhead = %0.2f"""', 'avg_sentences_per_second'], {}), "(\n    'Throughput Average (sentences/sec) with overhead = %0.2f',\n    avg_sentences_per_second)\n", (23906, 24001), True, 'import tensorflow as tf\n'), ((24001, 24101), 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', (['"""Throughput Average (sentences/sec) = %0.2f"""', 'ss_sentences_per_second'], {}), "('Throughput Average (sentences/sec) = %0.2f',\n    ss_sentences_per_second)\n", (24026, 24101), True, 'import tensorflow as tf\n'), ((24106, 24164), 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', (['"""-----------------------------"""'], {}), "('-----------------------------')\n", (24131, 24164), True, 'import tensorflow as tf\n'), ((24954, 24991), 'utils.utils.LogEvalRunHook', 'LogEvalRunHook', (['FLAGS.eval_batch_size'], {}), '(FLAGS.eval_batch_size)\n', (24968, 24991), False, 'from utils.utils import LogEvalRunHook, LogTrainRunHook, setup_xla_flags\n'), ((25126, 25137), 'time.time', 'time.time', ([], {}), '()\n', (25135, 25137), False, 'import time\n'), ((27372, 27412), 'tensorflow.io.gfile.GFile', 'tf.io.gfile.GFile', (['output_eval_file', '"""w"""'], {}), "(output_eval_file, 'w')\n", (27389, 27412), True, 'import tensorflow as tf\n'), ((27430, 27483), 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', (['"""***** Eval results *****"""'], {}), "('***** Eval results *****')\n", (27455, 27483), True, 'import tensorflow as tf\n'), ((28663, 28703), 'utils.utils.LogEvalRunHook', 'LogEvalRunHook', (['FLAGS.predict_batch_size'], {}), '(FLAGS.predict_batch_size)\n', (28677, 28703), False, 'from utils.utils import LogEvalRunHook, LogTrainRunHook, setup_xla_flags\n'), ((28829, 28872), 'tensorflow.io.gfile.GFile', 'tf.io.gfile.GFile', (['output_predict_file', '"""w"""'], {}), "(output_predict_file, 'w')\n", (28846, 28872), True, 'import tensorflow as tf\n'), ((28892, 28948), 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', (['"""***** Predict results *****"""'], {}), "('***** Predict results *****')\n", (28917, 28948), True, 'import tensorflow as tf\n'), ((29302, 29313), 'time.time', 'time.time', ([], {}), '()\n', (29311, 29313), False, 'import time\n'), ((5470, 5484), 'tensorflow.to_int32', 'tf.to_int32', (['t'], {}), '(t)\n', (5481, 5484), True, 'import tensorflow as tf\n'), ((9141, 9233), 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', (['"""  name = %s, shape = %s%s"""', 'var.name', 'var.shape', 'init_string'], {}), "('  name = %s, shape = %s%s', var.name, var.shape,\n    init_string)\n", (9166, 9233), True, 'import tensorflow as tf\n'), ((10846, 10915), 'tensorflow.metrics.false_negatives', 'tf.metrics.false_negatives', ([], {'labels': 'label_ids', 'predictions': 'predictions'}), '(labels=label_ids, predictions=predictions)\n', (10872, 10915), True, 'import tensorflow as tf\n'), ((10940, 11009), 'tensorflow.metrics.false_positives', 'tf.metrics.false_positives', ([], {'labels': 'label_ids', 'predictions': 'predictions'}), '(labels=label_ids, predictions=predictions)\n', (10966, 11009), True, 'import tensorflow as tf\n'), ((11034, 11102), 'tensorflow.metrics.true_positives', 'tf.metrics.true_positives', ([], {'labels': 'label_ids', 'predictions': 'predictions'}), '(labels=label_ids, predictions=predictions)\n', (11059, 11102), True, 'import tensorflow as tf\n'), ((11127, 11195), 'tensorflow.metrics.true_negatives', 'tf.metrics.true_negatives', ([], {'labels': 'label_ids', 'predictions': 'predictions'}), '(labels=label_ids, predictions=predictions)\n', (11152, 11195), True, 'import tensorflow as tf\n'), ((13290, 13352), 'tensorflow.estimator.EstimatorSpec', 'tf.estimator.EstimatorSpec', ([], {'mode': 'mode', 'predictions': 'predictions'}), '(mode=mode, predictions=predictions)\n', (13316, 13352), True, 'import tensorflow as tf\n'), ((14478, 14570), 'tensorflow.compat.v1.logging.info', 'tf.compat.v1.logging.info', (['"""  name = %s, shape = %s%s"""', 'var.name', 'var.shape', 'init_string'], {}), "('  name = %s, shape = %s%s', var.name, var.shape,\n    init_string)\n", (14503, 14570), True, 'import tensorflow as tf\n'), ((15054, 15064), 'tensorflow.no_op', 'tf.no_op', ([], {}), '()\n', (15062, 15064), True, 'import tensorflow as tf\n'), ((15472, 15564), 'tensorflow.estimator.EstimatorSpec', 'tf.estimator.EstimatorSpec', ([], {'mode': 'mode', 'loss': 'total_loss', 'eval_metric_ops': 'eval_metric_ops'}), '(mode=mode, loss=total_loss, eval_metric_ops=\n    eval_metric_ops)\n', (15498, 15564), True, 'import tensorflow as tf\n'), ((15618, 15628), 'tensorflow.no_op', 'tf.no_op', ([], {}), '()\n', (15626, 15628), True, 'import tensorflow as tf\n'), ((15890, 15954), 'tensorflow.estimator.EstimatorSpec', 'tf.estimator.EstimatorSpec', ([], {'mode': 'mode', 'predictions': 'probabilities'}), '(mode=mode, predictions=probabilities)\n', (15916, 15954), True, 'import tensorflow as tf\n'), ((16970, 17046), 'tensorflow.constant', 'tf.constant', (['all_input_ids'], {'shape': '[num_examples, seq_length]', 'dtype': 'tf.int32'}), '(all_input_ids, shape=[num_examples, seq_length], dtype=tf.int32)\n', (16981, 17046), True, 'import tensorflow as tf\n'), ((17115, 17192), 'tensorflow.constant', 'tf.constant', (['all_input_mask'], {'shape': '[num_examples, seq_length]', 'dtype': 'tf.int32'}), '(all_input_mask, shape=[num_examples, seq_length], dtype=tf.int32)\n', (17126, 17192), True, 'import tensorflow as tf\n'), ((17278, 17356), 'tensorflow.constant', 'tf.constant', (['all_segment_ids'], {'shape': '[num_examples, seq_length]', 'dtype': 'tf.int32'}), '(all_segment_ids, shape=[num_examples, seq_length], dtype=tf.int32)\n', (17289, 17356), True, 'import tensorflow as tf\n'), ((17440, 17504), 'tensorflow.constant', 'tf.constant', (['all_label_ids'], {'shape': '[num_examples]', 'dtype': 'tf.int32'}), '(all_label_ids, shape=[num_examples], dtype=tf.int32)\n', (17451, 17504), True, 'import tensorflow as tf\n'), ((19673, 19708), 'horovod.tensorflow.BroadcastGlobalVariablesHook', 'hvd.BroadcastGlobalVariablesHook', (['(0)'], {}), '(0)\n', (19705, 19708), True, 'import horovod.tensorflow as hvd\n'), ((21309, 21319), 'horovod.tensorflow.size', 'hvd.size', ([], {}), '()\n', (21317, 21319), True, 'import horovod.tensorflow as hvd\n'), ((21360, 21370), 'horovod.tensorflow.size', 'hvd.size', ([], {}), '()\n', (21368, 21370), True, 'import horovod.tensorflow as hvd\n'), ((21380, 21390), 'horovod.tensorflow.rank', 'hvd.rank', ([], {}), '()\n', (21388, 21390), True, 'import horovod.tensorflow as hvd\n'), ((5831, 5841), 'horovod.tensorflow.size', 'hvd.size', ([], {}), '()\n', (5839, 5841), True, 'import horovod.tensorflow as hvd\n'), ((5843, 5853), 'horovod.tensorflow.rank', 'hvd.rank', ([], {}), '()\n', (5851, 5853), True, 'import horovod.tensorflow as hvd\n'), ((11350, 11378), 'tensorflow.identity', 'tf.identity', (['MCC'], {'name': '"""MCC"""'}), "(MCC, name='MCC')\n", (11361, 11378), True, 'import tensorflow as tf\n'), ((11479, 11541), 'tensorflow.metrics.accuracy', 'tf.metrics.accuracy', ([], {'labels': 'label_ids', 'predictions': 'predictions'}), '(labels=label_ids, predictions=predictions)\n', (11498, 11541), True, 'import tensorflow as tf\n'), ((11578, 11618), 'tensorflow.metrics.mean', 'tf.metrics.mean', ([], {'values': 'per_example_loss'}), '(values=per_example_loss)\n', (11593, 11618), True, 'import tensorflow as tf\n'), ((11636, 11728), 'tf_metrics.f1', 'tf_metrics.f1', ([], {'labels': 'label_ids', 'predictions': 'predictions', 'num_classes': '(2)', 'pos_indices': '[1]'}), '(labels=label_ids, predictions=predictions, num_classes=2,\n    pos_indices=[1])\n', (11649, 11728), False, 'import tf_metrics\n'), ((11906, 11968), 'tensorflow.metrics.accuracy', 'tf.metrics.accuracy', ([], {'labels': 'label_ids', 'predictions': 'predictions'}), '(labels=label_ids, predictions=predictions)\n', (11925, 11968), True, 'import tensorflow as tf\n'), ((12005, 12045), 'tensorflow.metrics.mean', 'tf.metrics.mean', ([], {'values': 'per_example_loss'}), '(values=per_example_loss)\n', (12020, 12045), True, 'import tensorflow as tf\n'), ((13522, 13614), 'tensorflow.estimator.EstimatorSpec', 'tf.estimator.EstimatorSpec', ([], {'mode': 'mode', 'loss': 'total_loss', 'eval_metric_ops': 'eval_metric_ops'}), '(mode=mode, loss=total_loss, eval_metric_ops=\n    eval_metric_ops)\n', (13548, 13614), True, 'import tensorflow as tf\n'), ((14002, 14012), 'horovod.tensorflow.rank', 'hvd.rank', ([], {}), '()\n', (14010, 14012), True, 'import horovod.tensorflow as hvd\n'), ((15190, 15229), 'tensorflow.train.experimental.FixedLossScale', 'tf.train.experimental.FixedLossScale', (['(1)'], {}), '(1)\n', (15226, 15229), True, 'import tensorflow as tf\n'), ((17572, 17582), 'horovod.tensorflow.size', 'hvd.size', ([], {}), '()\n', (17580, 17582), True, 'import horovod.tensorflow as hvd\n'), ((17584, 17594), 'horovod.tensorflow.rank', 'hvd.rank', ([], {}), '()\n', (17592, 17594), True, 'import horovod.tensorflow as hvd\n'), ((21426, 21436), 'horovod.tensorflow.rank', 'hvd.rank', ([], {}), '()\n', (21434, 21436), True, 'import horovod.tensorflow as hvd\n'), ((21921, 21931), 'horovod.tensorflow.size', 'hvd.size', ([], {}), '()\n', (21929, 21931), True, 'import horovod.tensorflow as hvd\n'), ((15321, 15366), 'optimization.LAMBOptimizer', 'optimization.LAMBOptimizer', ([], {'learning_rate': '(0.0)'}), '(learning_rate=0.0)\n', (15347, 15366), False, 'import optimization\n'), ((15823, 15868), 'optimization.LAMBOptimizer', 'optimization.LAMBOptimizer', ([], {'learning_rate': '(0.0)'}), '(learning_rate=0.0)\n', (15849, 15868), False, 'import optimization\n'), ((21243, 21253), 'horovod.tensorflow.size', 'hvd.size', ([], {}), '()\n', (21251, 21253), True, 'import horovod.tensorflow as hvd\n'), ((21559, 21569), 'horovod.tensorflow.rank', 'hvd.rank', ([], {}), '()\n', (21567, 21569), True, 'import horovod.tensorflow as hvd\n')]