text: # Copyright 2018 The Texar Authors. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #      http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  """Interpolation Algorithm.
  """
  from __future__ import absolute_import
  from __future__ import print_function
  from __future__ import division
  from __future__ import unicode_literals
  
  import importlib
  from io import open
  
  import tensorflow as tf
  import texar as tx
  import numpy as np
  
  from interpolation_decoder import InterpolationDecoder
  from interpolation_helper import InterpolationHelper
  from rouge import Rouge
  
  flags = tf.flags
  
  flags.DEFINE_string("config_model", "configs.config_model", "The model config.")
  flags.DEFINE_string("config_data", "configs.config_iwslt14",
                      "The dataset config.")
  
  flags.DEFINE_string('lambdas_init', '[0.04,0.96,0.0]',
                      'initial value of lambdas')
  
  flags.DEFINE_float('delta_lambda_reward', 0.06,
                     'increment of lambda_reward every annealing')
  flags.DEFINE_float('delta_lambda_self', 0.06,
                     'decrement of lambda_self every annealing')
  flags.DEFINE_integer('lambda_reward_steps', 4,
                       'times of increasing lambda_reward '
                       'after incresing lambda_self once')
  
  flags.DEFINE_string('output_dir', '.', 'where to keep training logs')
  
  FLAGS = flags.FLAGS
  
  config_model = importlib.import_module(FLAGS.config_model)
  config_data = importlib.import_module(FLAGS.config_data)
  
  FLAGS.lambdas_init = eval(FLAGS.lambdas_init)
  
  if not FLAGS.output_dir.endswith('/'):
      FLAGS.output_dir += '/'
  log_dir = FLAGS.output_dir + 'training_log_interpolation' +\
            '_init' + '_' + str(FLAGS.lambdas_init[0]) +\
            '_' + str(FLAGS.lambdas_init[1]) +\
            '_' + str(FLAGS.lambdas_init[2]) +\
            '_dr' + str(FLAGS.delta_lambda_reward) +\
            '_ds' + str(FLAGS.delta_lambda_self) +\
            '_rstep' + str(FLAGS.lambda_reward_steps) + '/'
  tx.utils.maybe_create_dir(log_dir)
  
  
  def build_model(batch, train_data, lambdas):
      """
      This function is basically the same as build_model() in
      baseline_seq2seq_attn.py, except the
      InterpolateDecoder and InterpolateHelper.
      """
      batch_size = tf.shape(batch['target_length'])[0]
  
      source_embedder = tx.modules.WordEmbedder(
          vocab_size=train_data.source_vocab.size, hparams=config_model.embedder)
  
      encoder = tx.modules.BidirectionalRNNEncoder(
          hparams=config_model.encoder)
  
      enc_outputs, _ = encoder(source_embedder(batch['source_text_ids']))
  
      target_embedder = tx.modules.WordEmbedder(
          vocab_size=train_data.target_vocab.size, hparams=config_model.embedder)
  
      decoder = InterpolationDecoder(
          memory=tf.concat(enc_outputs, axis=2),
          memory_sequence_length=batch['source_length'],
          vocab_size=train_data.target_vocab.size,
          hparams=config_model.decoder)
  
      start_tokens = tf.ones_like(
          batch['target_length']) * train_data.target_vocab.bos_token_id
      helper = InterpolationHelper(
          embedding=target_embedder,
          start_tokens=start_tokens,
          end_token=train_data.target_vocab.eos_token_id,
          reward_metric=config_data.eval_metric,
          vocab=train_data.target_vocab,
          ground_truth=batch['target_text_ids'][:, 1:],
          ground_truth_length=batch['target_length'] - 1,
          lambdas=lambdas,)
  
      training_outputs, _, training_length = decoder(
          helper=helper,
          initial_state=decoder.zero_state(
              batch_size=batch_size, dtype=tf.float32),
          max_decoding_length=60)
  
      train_op = tx.core.get_train_op(
          tx.losses.sequence_sparse_softmax_cross_entropy(
              labels=training_outputs.sample_id,
              logits=training_outputs.logits,
              sequence_length=training_length),
          hparams=config_model.opt)
  
      beam_search_outputs, _, _ = \
          tx.modules.beam_search_decode(
              decoder_or_cell=decoder,
              embedding=target_embedder,
              start_tokens=start_tokens,
              end_token=train_data.target_vocab.eos_token_id,
              beam_width=config_model.beam_width,
              max_decoding_length=60)
  
      return train_op, beam_search_outputs
  
  
  def print_stdout_and_file(content, file):
      print(content)
      print(content, file=file)
  
  
  def main():
      """Entrypoint.
      """
      training_data = tx.data.PairedTextData(hparams=config_data.train)
      val_data = tx.data.PairedTextData(hparams=config_data.val)
      test_data = tx.data.PairedTextData(hparams=config_data.test)
      data_iterator = tx.data.TrainTestDataIterator(
          train=training_data, val=val_data, test=test_data)
  
      batch = data_iterator.get_next()
      lambdas_ts = tf.placeholder(shape=[3], dtype=tf.float32)
  
      train_op, infer_outputs = build_model(batch, training_data, lambdas_ts)
  
      def _train_epoch(sess, epoch, lambdas):
          data_iterator.switch_to_train_data(sess)
          log_file = open(log_dir + 'training_log' + str(epoch) + '.txt', 'w',
                          encoding='utf-8')
  
          step = 0
          while True:
              try:
                  loss = sess.run(train_op, feed_dict={
                      lambdas_ts: np.array(lambdas)})
                  print("step={}, loss={:.4f}, lambdas={}".format(
                      step, loss, lambdas), file=log_file)
                  if step % config_data.observe_steps == 0:
                      print("step={}, loss={:.4f}, lambdas={}".format(
                          step, loss, lambdas))
                  log_file.flush()
                  step += 1
  
              except tf.errors.OutOfRangeError:
                  break
  
      def _eval_epoch(sess, mode, epoch_no):
          """
          This function is the same as _eval_epoch() in
          baseline_seq2seq_attn_main.py.
          """
          if mode == 'val':
              data_iterator.switch_to_val_data(sess)
          else:
              data_iterator.switch_to_test_data(sess)
  
          refs, hypos = [], []
          while True:
              try:
                  fetches = [
                      batch['target_text'][:, 1:],
                      infer_outputs.predicted_ids[:, :, 0]
                  ]
                  feed_dict = {
                      tx.global_mode(): tf.estimator.ModeKeys.EVAL
                  }
                  target_texts_ori, output_ids = \
                      sess.run(fetches, feed_dict=feed_dict)
  
                  target_texts = tx.utils.strip_special_tokens(
                      target_texts_ori.tolist(), is_token_list=True)
                  target_texts = tx.utils.str_join(target_texts)
                  output_texts = tx.utils.map_ids_to_strs(
                      ids=output_ids, vocab=val_data.target_vocab)
  
                  tx.utils.write_paired_text(
                      target_texts, output_texts,
                      log_dir + mode + '_results' + str(epoch_no) + '.txt',
                      append=True, mode='h', sep=' ||| ')
  
                  for hypo, ref in zip(output_texts, target_texts):
                      if config_data.eval_metric == 'bleu':
                          hypos.append(hypo)
                          refs.append([ref])
                      elif config_data.eval_metric == 'rouge':
                          hypos.append(tx.utils.compat_as_text(hypo))
                          refs.append(tx.utils.compat_as_text(ref))
              except tf.errors.OutOfRangeError:
                  break
  
          if config_data.eval_metric == 'bleu':
              return tx.evals.corpus_bleu_moses(
                  list_of_references=refs, hypotheses=hypos)
          elif config_data.eval_metric == 'rouge':
              rouge = Rouge()
              return rouge.get_scores(hyps=hypos, refs=refs, avg=True)
  
      def _calc_reward(score):
          """
          Return the bleu score or the sum of (Rouge-1, Rouge-2, Rouge-L).
          """
          if config_data.eval_metric == 'bleu':
              return score
          elif config_data.eval_metric == 'rouge':
              return sum([value['f'] for key, value in score.items()])
  
      def _anneal():
          """
          Operate lambdas when the reward of val set decrease.
          """
          def _update_self():
              """
              Decrease lambda_truth and increase lambda_self.
              """
              lambdas[1] -= FLAGS.delta_lambda_self
              lambdas[0] += FLAGS.delta_lambda_self
              updates.append('self')
  
          def _update_rew():
              """
              Decrease lambda_truth and increase lambda_reward.
              """
              lambdas[1] -= FLAGS.delta_lambda_reward
              lambdas[2] += FLAGS.delta_lambda_reward
              updates.append('rew')
  
          if updates[-FLAGS.lambda_reward_steps:] == \
                  ['rew'] * FLAGS.lambda_reward_steps:
              _update_self()
          else:
              _update_rew()
  
      saver = tf.train.Saver(max_to_keep=2)
      with tf.Session() as sess:
          sess.run(tf.global_variables_initializer())
          sess.run(tf.local_variables_initializer())
          sess.run(tf.tables_initializer())
  
          lambdas = FLAGS.lambdas_init
          updates = ['rew'] * FLAGS.lambda_reward_steps
  
          best_val_score, best_val_score_current_lambdas = -1., -1.
          scores_file = open(log_dir + 'scores.txt', 'w', encoding='utf-8')
  
          for i in range(config_data.num_epochs):
              print_stdout_and_file(
                  'training epoch={}, lambdas={}'.format(i, lambdas),
                  file=scores_file)
              _train_epoch(sess, i, lambdas)
              saver.save(sess, log_dir + 'models/model{}.ckpt'.format(i))
  
              val_score = _eval_epoch(sess, 'val', i)
              test_score = _eval_epoch(sess, 'test', i)
  
              if _calc_reward(val_score) < best_val_score_current_lambdas:
                  _anneal()
                  best_val_score_current_lambdas = -1.
                  saver.restore(
                      sess, log_dir + 'models/model{}.ckpt'.format(i - 1))
              else:
                  best_val_score_current_lambdas = _calc_reward(val_score)
  
              best_val_score = max(best_val_score, _calc_reward(val_score))
  
              if config_data.eval_metric == 'bleu':
                  print_stdout_and_file(
                      'val epoch={}, BLEU={:.4f}; best-ever={:.4f}'.format(
                          i, val_score, best_val_score), file=scores_file)
  
                  print_stdout_and_file(
                      'test epoch={}, BLEU={:.4f}'.format(i, test_score),
                      file=scores_file)
                  print_stdout_and_file('=' * 50, file=scores_file)
  
              elif config_data.eval_metric == 'rouge':
                  print_stdout_and_file(
                      'valid epoch {}:'.format(i), file=scores_file)
                  for key, value in val_score.items():
                      print_stdout_and_file(
                          '{}: {}'.format(key, value), file=scores_file)
                  print_stdout_and_file('fsum: {}; best_val_fsum: {}'.format(
                      _calc_reward(val_score), best_val_score), file=scores_file)
  
                  print_stdout_and_file(
                      'test epoch {}:'.format(i), file=scores_file)
                  for key, value in test_score.items():
                      print_stdout_and_file(
                          '{}: {}'.format(key, value), file=scores_file)
                  print_stdout_and_file('=' * 110, file=scores_file)
  
              scores_file.flush()
  
  
  if __name__ == '__main__':
      main()
  