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.
  """Hierarchical Recurrent Encoder-Decoder (HRED) for dialog response
  generation.
  """
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  # pylint: disable=invalid-name, too-many-locals
  
  import importlib
  import numpy as np
  import tensorflow as tf
  import texar as tx
  
  from nltk.translate.bleu_score import sentence_bleu
  from nltk.translate.bleu_score import SmoothingFunction
  
  flags = tf.flags
  
  flags.DEFINE_string('config_data', 'config_data', 'The data config')
  flags.DEFINE_string('config_model', 'config_model_biminor', 'The model config')
  
  FLAGS = flags.FLAGS
  
  config_data = importlib.import_module(FLAGS.config_data)
  config_model = importlib.import_module(FLAGS.config_model)
  
  encoder_hparams = config_model.encoder_hparams
  decoder_hparams = config_model.decoder_hparams
  opt_hparams = config_model.opt_hparams
  
  def main():
      """Entrypoint.
      """
      # Data
      train_data = tx.data.MultiAlignedData(config_data.data_hparams['train'])
      val_data = tx.data.MultiAlignedData(config_data.data_hparams['val'])
      test_data = tx.data.MultiAlignedData(config_data.data_hparams['test'])
      iterator = tx.data.TrainTestDataIterator(train=train_data,
                                               val=val_data,
                                               test=test_data)
      data_batch = iterator.get_next()
  
      # (speaker's meta info)
      spk_src = tf.stack([data_batch['spk_{}'.format(i)]
                          for i in range(config_data.max_utterance_cnt)], 1)
      spk_tgt = data_batch['spk_tgt']
      def _add_source_speaker_token(x):
          return tf.concat([x, tf.reshape(spk_src, (-1, 1))], 1)
      def _add_target_speaker_token(x):
          return (x, ) + (tf.reshape(spk_tgt, (-1, 1)), )
  
      # HRED model
      embedder = tx.modules.WordEmbedder(
          init_value=train_data.embedding_init_value(0).word_vecs)
      encoder = tx.modules.HierarchicalRNNEncoder(hparams=encoder_hparams)
  
      decoder = tx.modules.BasicRNNDecoder(
          hparams=decoder_hparams, vocab_size=train_data.vocab(0).size)
  
      connector = tx.modules.connectors.MLPTransformConnector(
          decoder.cell.state_size)
  
      context_embed = embedder(data_batch['source_text_ids'])
      ecdr_states = encoder(
          context_embed,
          medium=['flatten', _add_source_speaker_token],
          sequence_length_minor=data_batch['source_length'],
          sequence_length_major=data_batch['source_utterance_cnt'])
      ecdr_states = ecdr_states[1]
  
      ecdr_states = _add_target_speaker_token(ecdr_states)
      dcdr_states = connector(ecdr_states)
  
      # (decoding for training)
      target_embed = embedder(data_batch['target_text_ids'])
      outputs, _, lengths = decoder(
          initial_state=dcdr_states,
          inputs=target_embed,
          sequence_length=data_batch['target_length'] - 1)
  
      # Sentence level lld, for training
      mle_loss = tx.losses.sequence_sparse_softmax_cross_entropy(
          labels=data_batch['target_text_ids'][:, 1:],
          logits=outputs.logits,
          sequence_length=lengths)
      # Token level lld, for perplexity evaluation
      avg_mle_loss = tx.losses.sequence_sparse_softmax_cross_entropy(
          labels=data_batch['target_text_ids'][:, 1:],
          logits=outputs.logits,
          sequence_length=lengths,
          sum_over_timesteps=False,
          average_across_timesteps=True)
      perplexity = tf.exp(avg_mle_loss)
  
      global_step = tf.Variable(0, name='global_step', trainable=True)
      train_op = tx.core.get_train_op(
          mle_loss, global_step=global_step, hparams=opt_hparams)
  
      # Decoding
  
      target_bos_token_id = train_data.vocab(0).bos_token_id
      target_eos_token_id = train_data.vocab(0).eos_token_id
      start_tokens = \
          tf.ones_like(data_batch['target_length']) * target_bos_token_id
  
      # Random sample decoding
      decoding_strategy = 'infer_' + 'sample'
      infer_samples, lengths = [], []
      for _ in range(config_model.num_samples):
          infer_outputs_i, _, lengths_i = decoder(
              decoding_strategy=decoding_strategy,
              initial_state=dcdr_states,
              start_tokens=start_tokens,
              end_token=target_eos_token_id,
              embedding=embedder,
              max_decoding_length=50)
          infer_samples.append(
              tf.expand_dims(infer_outputs_i.sample_id, axis=2))
          lengths.append(tf.expand_dims(lengths_i, axis=1))
  
      infer_samples = tx.utils.pad_and_concat(
          infer_samples, axis=2, pad_axis=1)
      rand_sample_text = train_data.vocab(0).map_ids_to_tokens(infer_samples)
      rand_lengths = tf.concat(lengths, axis=1)
  
      # Beam search decoding
      beam_search_samples, beam_states, _ = tx.modules.beam_search_decode(
          decoder,
          initial_state=dcdr_states,
          start_tokens=start_tokens,
          end_token=target_eos_token_id,
          embedding=embedder,
          beam_width=config_model.beam_width,
          max_decoding_length=50)
  
      beam_sample_text = train_data.vocab(0).map_ids_to_tokens(
          beam_search_samples.predicted_ids)
      beam_lengths = beam_states.lengths
  
      # Running procedures
  
      def _train_epoch(sess, epoch, display=1000):
          iterator.switch_to_train_data(sess)
  
          while True:
              try:
                  feed = {tx.global_mode(): tf.estimator.ModeKeys.TRAIN}
                  step, loss, _ = sess.run(
                      [global_step, mle_loss, train_op], feed_dict=feed)
  
                  if step % display == 0:
                      print('step {} at epoch {}: loss={}'.format(
                          step, epoch, loss))
  
              except tf.errors.OutOfRangeError:
                  break
  
          print('epoch {} train: loss={}'.format(epoch, loss))
  
      def _test_epoch_ppl(sess, epoch):
          iterator.switch_to_test_data(sess)
  
          pples = []
          while True:
              try:
                  feed = {tx.global_mode(): tf.estimator.ModeKeys.EVAL}
                  ppl = sess.run(perplexity, feed_dict=feed)
                  pples.append(ppl)
  
              except tf.errors.OutOfRangeError:
                  avg_ppl = np.mean(pples)
                  print('epoch {} perplexity={}'.format(epoch, avg_ppl))
                  break
  
      def _test_epoch_bleu(sess, epoch, sample_text, sample_lengths):
          iterator.switch_to_test_data(sess)
  
          bleu_prec = [[] for i in range(1, 5)]
          bleu_recall = [[] for i in range(1, 5)]
  
          def _bleus(ref, sample):
              res = []
              for weight in [[1, 0, 0, 0],
                             [1, 0, 0, 0],
                             [1/2., 1/2., 0, 0],
                             [1/3., 1/3., 1/3., 0],
                             [1/4., 1/4., 1/4., 1/4.]]:
                  res.append(sentence_bleu(
                      [ref],
                      sample,
                      smoothing_function=SmoothingFunction().method7,
                      weights=weight))
              return res
  
          while True:
              try:
                  feed = {tx.global_mode(): tf.estimator.ModeKeys.EVAL}
  
                  samples_, sample_lengths_, references, refs_cnt = \
                      sess.run([sample_text,
                                sample_lengths,
                                data_batch['refs_text'][:, :, 1:],
                                data_batch['refs_utterance_cnt']],
                               feed_dict=feed)
  
                  samples_ = np.transpose(samples_, (0, 2, 1))
                  samples_ = [
                      [sample[:l] for sample, l in zip(beam, lens)]
                      for beam, lens in zip(samples_.tolist(), sample_lengths_)
                  ]
                  references = [
                      [ref[:ref.index(b'<EOS>')] for ref in refs[:cnt]]
                      for refs, cnt in zip(references.tolist(), refs_cnt)
                  ]
  
                  for beam, refs in zip(samples_, references):
                      bleu_scores = [
                          [_bleus(ref, sample) for ref in refs]
                          for sample in beam
                      ]
                      bleu_scores = np.transpose(np.array(bleu_scores), (2, 0, 1))
  
                      for i in range(1, 5):
                          bleu_i = bleu_scores[i]
                          bleu_i_precision = bleu_i.max(axis=1).mean()
                          bleu_i_recall = bleu_i.max(axis=0).mean()
  
                          bleu_prec[i-1].append(bleu_i_precision)
                          bleu_recall[i-1].append(bleu_i_recall)
  
  
              except tf.errors.OutOfRangeError:
                  break
  
          bleu_prec = [np.mean(x) for x in bleu_prec]
          bleu_recall = [np.mean(x) for x in bleu_recall]
  
          print('epoch {}:'.format(epoch))
          for i in range(1, 5):
              print(' -- bleu-{} prec={}, recall={}'.format(
                  i, bleu_prec[i-1], bleu_recall[i-1]))
  
      with tf.Session() as sess:
          sess.run(tf.global_variables_initializer())
          sess.run(tf.local_variables_initializer())
          sess.run(tf.tables_initializer())
  
          num_epochs = 10
          for epoch in range(1, num_epochs+1):
              _train_epoch(sess, epoch)
              _test_epoch_ppl(sess, epoch)
  
              if epoch % 5 == 0:
                  print('random sample: ')
                  _test_epoch_bleu(sess, epoch, rand_sample_text, rand_lengths)
                  print('beam-search: ')
                  _test_epoch_bleu(sess, epoch, beam_sample_text, beam_lengths)
  
          if num_epochs % 5 != 0:
              print('random sample: ')
              _test_epoch_bleu(sess, num_epochs, rand_sample_text, rand_lengths)
              print('beam-search: ')
              _test_epoch_bleu(sess, num_epochs, beam_sample_text, beam_lengths)
  
  
  if __name__ == "__main__":
      main()
  