text: #!/usr/bin/env python
  # coding: utf-8
  
  # In[1]:
  
  
  # A BERT model with history answer embedding (HAE)
  
  
  # In[2]:
  
  
  # import os
  # os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
  # os.environ["CUDA_VISIBLE_DEVICES"]="3"
  
  
  # In[3]:
  
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import collections
  import json
  import math
  import os
  import modeling
  import optimization
  import tokenization
  import six
  import tensorflow as tf
  import numpy as np
  from copy import deepcopy
  import pickle
  import itertools
  from time import time
  
  from cqa_supports import *
  from cqa_flags import FLAGS
  from cqa_model import *
  from cqa_gen_batches import *
  
  from scorer import external_call # quac official evaluation script
  
  
  # In[4]:
  
  
  for key in FLAGS:
      print(key, ':', FLAGS[key].value)
  
  tf.set_random_seed(0)
  tf.logging.set_verbosity(tf.logging.INFO)
  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.gfile.MakeDirs(FLAGS.output_dir)
  tf.gfile.MakeDirs(FLAGS.output_dir + '/summaries/train/')
  tf.gfile.MakeDirs(FLAGS.output_dir + '/summaries/val/')
  tf.gfile.MakeDirs(FLAGS.output_dir + '/summaries/rl/')
  
  tokenizer = tokenization.FullTokenizer(vocab_file=FLAGS.vocab_file, do_lower_case=FLAGS.do_lower_case)
  
  if FLAGS.do_train:
      # read in training data, generate training features, and generate training batches
      train_examples = None
      num_train_steps = None
      num_warmup_steps = None
      train_file = FLAGS.quac_train_file
      train_examples = read_quac_examples(input_file=train_file, is_training=True)
          
      
      # we attempt to read features from cache
      features_fname = FLAGS.cache_dir + FLAGS.dataset.lower() +                '/train_features_{}_{}.pkl'.format(FLAGS.load_small_portion, FLAGS.max_considered_history_turns)
      example_tracker_fname = FLAGS.cache_dir + FLAGS.dataset.lower() +                '/example_tracker_{}_{}.pkl'.format(FLAGS.load_small_portion, FLAGS.max_considered_history_turns)
      variation_tracker_fname = FLAGS.cache_dir + FLAGS.dataset.lower() +                '/variation_tracker_{}_{}.pkl'.format(FLAGS.load_small_portion, FLAGS.max_considered_history_turns)
      example_features_nums_fname = FLAGS.cache_dir + FLAGS.dataset.lower() +                '/example_features_nums_{}_{}.pkl'.format(FLAGS.load_small_portion, FLAGS.max_considered_history_turns)
          
      try:
          print('attempting to load train features from cache')
          with open(features_fname, 'rb') as handle:
              train_features = pickle.load(handle)
          with open(example_tracker_fname, 'rb') as handle:
              example_tracker = pickle.load(handle)
          with open(variation_tracker_fname, 'rb') as handle:
              variation_tracker = pickle.load(handle)
          with open(example_features_nums_fname, 'rb') as handle:
              example_features_nums = pickle.load(handle)
      except:
          print('train feature cache does not exist, generating')
          train_features, example_tracker, variation_tracker,                                 example_features_nums = convert_examples_to_variations_and_then_features(
                                          examples=train_examples, tokenizer=tokenizer, 
                                          max_seq_length=FLAGS.max_seq_length, doc_stride=FLAGS.doc_stride, 
                                          max_query_length=FLAGS.max_query_length, 
                                          max_considered_history_turns=FLAGS.max_considered_history_turns, 
                                          is_training=True)
          with open(features_fname, 'wb') as handle:
              pickle.dump(train_features, handle)
          with open(example_tracker_fname, 'wb') as handle:
              pickle.dump(example_tracker, handle)
          with open(variation_tracker_fname, 'wb') as handle:
              pickle.dump(variation_tracker, handle)     
          with open(example_features_nums_fname, 'wb') as handle:
              pickle.dump(example_features_nums, handle) 
          print('train features generated')
                  
      train_batches = cqa_gen_example_aware_batches(train_features, example_tracker, variation_tracker, 
                                                    example_features_nums, FLAGS.train_batch_size, 
                                                    FLAGS.num_train_epochs, shuffle=False)
      
      num_train_steps = FLAGS.train_steps
      num_warmup_steps = int(num_train_steps * FLAGS.warmup_proportion)
  
  if FLAGS.do_predict:
      # read in validation data, generate val features
      val_file = FLAGS.quac_predict_file
      val_examples = read_quac_examples(input_file=val_file, is_training=False)
      
      # we read in the val file in json for the external_call function in the validation step
      val_file_json = json.load(open(val_file, 'r'))['data']
      
      # we attempt to read features from cache
      features_fname = FLAGS.cache_dir + FLAGS.dataset.lower() +                      '/val_features_{}_{}.pkl'.format(FLAGS.load_small_portion, FLAGS.max_considered_history_turns)
      example_tracker_fname = FLAGS.cache_dir + FLAGS.dataset.lower() +                      '/val_example_tracker_{}_{}.pkl'.format(FLAGS.load_small_portion, FLAGS.max_considered_history_turns)
      variation_tracker_fname = FLAGS.cache_dir + FLAGS.dataset.lower() +                      '/val_variation_tracker_{}_{}.pkl'.format(FLAGS.load_small_portion, FLAGS.max_considered_history_turns)
      example_features_nums_fname = FLAGS.cache_dir + FLAGS.dataset.lower() +                      '/val_example_features_nums_{}_{}.pkl'.format(FLAGS.load_small_portion, FLAGS.max_considered_history_turns)
          
      try:
          print('attempting to load val features from cache')
          with open(features_fname, 'rb') as handle:
              val_features = pickle.load(handle)
          with open(example_tracker_fname, 'rb') as handle:
              val_example_tracker = pickle.load(handle)
          with open(variation_tracker_fname, 'rb') as handle:
              val_variation_tracker = pickle.load(handle)
          with open(example_features_nums_fname, 'rb') as handle:
              val_example_features_nums = pickle.load(handle)
      except:
          print('val feature cache does not exist, generating')
          val_features, val_example_tracker, val_variation_tracker, val_example_features_nums =                                                    convert_examples_to_variations_and_then_features(
                                                     examples=val_examples, tokenizer=tokenizer, 
                                                     max_seq_length=FLAGS.max_seq_length, doc_stride=FLAGS.doc_stride, 
                                                     max_query_length=FLAGS.max_query_length, 
                                                     max_considered_history_turns=FLAGS.max_considered_history_turns, 
                                                     is_training=False)
          with open(features_fname, 'wb') as handle:
              pickle.dump(val_features, handle)
          with open(example_tracker_fname, 'wb') as handle:
              pickle.dump(val_example_tracker, handle)
          with open(variation_tracker_fname, 'wb') as handle:
              pickle.dump(val_variation_tracker, handle)  
          with open(example_features_nums_fname, 'wb') as handle:
              pickle.dump(val_example_features_nums, handle)
          print('val features generated')
      
      
      num_val_examples = len(val_examples)
      
  
  # tf Graph input
  unique_ids = tf.placeholder(tf.int32, shape=[None], name='unique_ids')
  input_ids = tf.placeholder(tf.int32, shape=[None, FLAGS.max_seq_length], name='input_ids')
  input_mask = tf.placeholder(tf.int32, shape=[None, FLAGS.max_seq_length], name='input_mask')
  segment_ids = tf.placeholder(tf.int32, shape=[None, FLAGS.max_seq_length], name='segment_ids')
  start_positions = tf.placeholder(tf.int32, shape=[None], name='start_positions')
  end_positions = tf.placeholder(tf.int32, shape=[None], name='end_positions')
  history_answer_marker = tf.placeholder(tf.int32, shape=[None, FLAGS.max_seq_length], name='history_answer_marker')
  training = tf.placeholder(tf.bool, name='training')
  get_segment_rep = tf.placeholder(tf.bool, name='get_segment_rep')
  
  
  bert_representation = bert_rep(
      bert_config=bert_config,
      is_training=training,
      input_ids=input_ids,
      input_mask=input_mask,
      segment_ids=segment_ids,
      history_answer_marker=history_answer_marker,
      use_one_hot_embeddings=False
      )
      
  (start_logits, end_logits) = cqa_model(bert_representation)
  
  
  tvars = tf.trainable_variables()
  
  initialized_variable_names = {}
  if FLAGS.init_checkpoint:
      (assignment_map, initialized_variable_names) = modeling.get_assigment_map_from_checkpoint(tvars, 
                                                                                                FLAGS.init_checkpoint)
      tf.train.init_from_checkpoint(FLAGS.init_checkpoint, assignment_map)
  
  # compute loss
  seq_length = modeling.get_shape_list(input_ids)[1]
  def compute_loss(logits, positions):
      one_hot_positions = tf.one_hot(
          positions, depth=seq_length, dtype=tf.float32)
      log_probs = tf.nn.log_softmax(logits, axis=-1)
      loss = -tf.reduce_mean(tf.reduce_sum(one_hot_positions * log_probs, axis=-1))
      return loss
  
  # get the max prob for the predicted start/end position
  start_probs = tf.nn.softmax(start_logits, axis=-1)
  start_prob = tf.reduce_max(start_probs, axis=-1)
  end_probs = tf.nn.softmax(end_logits, axis=-1)
  end_prob = tf.reduce_max(end_probs, axis=-1)
  
  start_loss = compute_loss(start_logits, start_positions)
  end_loss = compute_loss(end_logits, end_positions)
  total_loss = (start_loss + end_loss) / 2.0
  tf.summary.scalar('total_loss', total_loss)
  
  
  if FLAGS.do_train:
      train_op = optimization.create_optimizer(total_loss, FLAGS.learning_rate, num_train_steps, num_warmup_steps, False)
  
      print("***** Running training *****")
      print("  Num orig examples = %d", len(train_examples))
      print("  Num train_features = %d", len(train_features))
      print("  Batch size = %d", FLAGS.train_batch_size)
      print("  Num steps = %d", num_train_steps)
      
  merged_summary_op = tf.summary.merge_all()
  
  RawResult = collections.namedtuple("RawResult", ["unique_id", "start_logits", "end_logits"])
  
  saver = tf.train.Saver()
  # Initializing the variables
  init = tf.global_variables_initializer()
  tf.get_default_graph().finalize()
  with tf.Session() as sess:
      sess.run(init)
  
      if FLAGS.do_train:
          train_summary_writer = tf.summary.FileWriter(FLAGS.output_dir + 'summaries/train', sess.graph)
          val_summary_writer = tf.summary.FileWriter(FLAGS.output_dir + 'summaries/val')
          
          f1_list = []
          heq_list = []
          dheq_list = []
          
          # Training cycle
          for step, batch in enumerate(train_batches):
              if step > num_train_steps:
                  # this means the learning rate has been decayed to 0
                  break
                  
              batch_features, batch_example_tracker, batch_variation_tracker = batch
              
  
              selected_example_features, relative_selected_pos = get_selected_example_features_without_actions(
                                                      batch_features, batch_example_tracker, batch_variation_tracker)
  
              fd = convert_features_to_feed_dict(selected_example_features) # feed_dict
              try:
                  _, train_summary, total_loss_res = sess.run([train_op, merged_summary_op, total_loss], 
                                             feed_dict={unique_ids: fd['unique_ids'], input_ids: fd['input_ids'], 
                                             input_mask: fd['input_mask'], segment_ids: fd['segment_ids'], 
                                             start_positions: fd['start_positions'], end_positions: fd['end_positions'], 
                                             history_answer_marker: fd['history_answer_marker'], training: True})
              except:
                  print('features length: ', len(selected_example_features))
  
              train_summary_writer.add_summary(train_summary, step)
              train_summary_writer.flush()
              print('training step: {}, total_loss: {}'.format(step, total_loss_res))
              
              if step >= FLAGS.evaluate_after and step % FLAGS.evaluation_steps == 0 and step != 0:
                  val_total_loss = []
                  all_results = []
                  all_selected_examples = []
                  all_selected_features = []
                  
                  total_num_selected = 0
                  total_num_actions = 0
                  total_num_examples = 0
                  
                  val_batches = cqa_gen_example_aware_batches(val_features, val_example_tracker, val_variation_tracker, 
                                             val_example_features_nums, FLAGS.predict_batch_size, 1, shuffle=False)
                  
                  for val_batch in val_batches:
  
                      batch_results = []
                      batch_features, batch_example_tracker, batch_variation_tracker = val_batch
                      
                      selected_example_features, relative_selected_pos = get_selected_example_features_without_actions(
                                                      batch_features, batch_example_tracker, batch_variation_tracker)
  
                          
                      try:
                          all_selected_features.extend(selected_example_features)
  
                          fd = convert_features_to_feed_dict(selected_example_features) # feed_dict
                          start_logits_res, end_logits_res, batch_total_loss = sess.run([start_logits, end_logits, total_loss], 
                                      feed_dict={unique_ids: fd['unique_ids'], input_ids: fd['input_ids'], 
                                      input_mask: fd['input_mask'], segment_ids: fd['segment_ids'], 
                                      start_positions: fd['start_positions'], end_positions: fd['end_positions'], 
                                      history_answer_marker: fd['history_answer_marker'], training: False})
  
                          val_total_loss.append(batch_total_loss)
  
                          for each_unique_id, each_start_logits, each_end_logits in zip(fd['unique_ids'], start_logits_res, 
                                                                                        end_logits_res):  
                              each_unique_id = int(each_unique_id)
                              each_start_logits = [float(x) for x in each_start_logits.flat]
                              each_end_logits = [float(x) for x in each_end_logits.flat]
                              batch_results.append(RawResult(unique_id=each_unique_id, start_logits=each_start_logits, 
                                                             end_logits=each_end_logits))
  
                          all_results.extend(batch_results)
                      except:
                          print('batch dropped because too large!')
  
                  output_prediction_file = os.path.join(FLAGS.output_dir, "predictions_{}.json".format(step))
                  output_nbest_file = os.path.join(FLAGS.output_dir, "nbest_predictions_{}.json".format(step))
  
                  write_predictions(val_examples, all_selected_features, all_results,
                                    FLAGS.n_best_size, FLAGS.max_answer_length,
                                    FLAGS.do_lower_case, output_prediction_file,
                                    output_nbest_file)
  
                  val_total_loss_value = np.average(val_total_loss)
                                  
                  
                  # call the official evaluation script
                  val_summary = tf.Summary() 
                  val_eval_res = external_call(val_file_json, output_prediction_file)
  
                  val_f1 = val_eval_res['f1']
                  val_followup = val_eval_res['followup']
                  val_yesno = val_eval_res['yes/no']
                  val_heq = val_eval_res['HEQ']
                  val_dheq = val_eval_res['DHEQ']
  
                  heq_list.append(val_heq)
                  dheq_list.append(val_dheq)
  
                  val_summary.value.add(tag="followup", simple_value=val_followup)
                  val_summary.value.add(tag="val_yesno", simple_value=val_yesno)
                  val_summary.value.add(tag="val_heq", simple_value=val_heq)
                  val_summary.value.add(tag="val_dheq", simple_value=val_dheq)
  
                  print('evaluation: {}, total_loss: {}, f1: {}, followup: {}, yesno: {}, heq: {}, dheq: {}\n'.format(
                      step, val_total_loss_value, val_f1, val_followup, val_yesno, val_heq, val_dheq))
                  with open(FLAGS.output_dir + 'step_result.txt', 'a') as fout:
                          fout.write('{},{},{},{},{},{}\n'.format(step, val_f1, val_heq, val_dheq, 
                                              FLAGS.history, FLAGS.output_dir))
                  
                  val_summary.value.add(tag="total_loss", simple_value=val_total_loss_value)
                  val_summary.value.add(tag="f1", simple_value=val_f1)
                  f1_list.append(val_f1)
                  
                  val_summary_writer.add_summary(val_summary, step)
                  val_summary_writer.flush()
                  
                  save_path = saver.save(sess, '{}/model_{}.ckpt'.format(FLAGS.output_dir, step))
                  print('Model saved in path', save_path)
  
                  
  
  
  # In[5]:
  
  
  best_f1 = max(f1_list)
  best_f1_idx = f1_list.index(best_f1)
  best_heq = heq_list[best_f1_idx]
  best_dheq = dheq_list[best_f1_idx]
  with open(FLAGS.output_dir + 'result.txt', 'w') as fout:
      fout.write('{},{},{},{},{}\n'.format(best_f1, best_heq, best_dheq, FLAGS.history, FLAGS.output_dir))
  
  
  # In[ ]:
  
  
  
  
  