text: # Copyright 2015 Google Inc. 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.
  # ==============================================================================
  """Neural GPU."""
  
  from __future__ import print_function
  
  import math
  import os
  import random
  import sys
  import threading
  import time
  
  import numpy as np
  from six.moves import xrange
  import tensorflow as tf
  
  import program_utils
  import data_utils as data
  import neural_gpu as ngpu
  import wmt_utils as wmt
  
  tf.app.flags.DEFINE_float("lr", 0.1, "Learning rate.")
  tf.app.flags.DEFINE_float("init_weight", 0.8, "Initial weights deviation.")
  tf.app.flags.DEFINE_float("max_grad_norm", 4.0, "Clip gradients to this norm.")
  tf.app.flags.DEFINE_float("cutoff", 1.2, "Cutoff at the gates.")
  tf.app.flags.DEFINE_float("curriculum_ppx", 9.9, "Move curriculum if ppl < X.")
  tf.app.flags.DEFINE_float("curriculum_seq", 0.3, "Move curriculum if seq < X.")
  tf.app.flags.DEFINE_float("dropout", 0.1, "Dropout that much.")
  tf.app.flags.DEFINE_float("grad_noise_scale", 0.0, "Gradient noise scale.")
  tf.app.flags.DEFINE_float("max_sampling_rate", 0.1, "Maximal sampling rate.")
  tf.app.flags.DEFINE_float("length_norm", 0.0, "Length normalization.")
  tf.app.flags.DEFINE_float("train_beam_freq", 0.0, "Beam-based training.")
  tf.app.flags.DEFINE_float("train_beam_anneal", 20000, "How many steps anneal.")
  tf.app.flags.DEFINE_integer("eval_beam_steps", 4, "How many beam steps eval.")
  tf.app.flags.DEFINE_integer("batch_size", 32, "Batch size.")
  tf.app.flags.DEFINE_integer("steps_per_checkpoint", 100, "Steps per epoch.")
  tf.app.flags.DEFINE_integer("nmaps", 64, "Number of floats in each cell.")
  tf.app.flags.DEFINE_integer("vec_size", 64, "Size of word vectors.")
  tf.app.flags.DEFINE_integer("train_data_size", 1000, "Training examples/len.")
  tf.app.flags.DEFINE_integer("max_length", 40, "Maximum length.")
  tf.app.flags.DEFINE_integer("random_seed", 125459, "Random seed.")
  tf.app.flags.DEFINE_integer("nconvs", 2, "How many convolutions / 1 step.")
  tf.app.flags.DEFINE_integer("kw", 3, "Kernel width.")
  tf.app.flags.DEFINE_integer("kh", 3, "Kernel height.")
  tf.app.flags.DEFINE_integer("height", 4, "Height.")
  tf.app.flags.DEFINE_integer("mem_size", -1, "Memory size (sqrt)")
  tf.app.flags.DEFINE_integer("soft_mem_size", 1024, "Softmax memory this size.")
  tf.app.flags.DEFINE_integer("num_gpus", 1, "Number of GPUs to use.")
  tf.app.flags.DEFINE_integer("num_replicas", 1, "Number of replicas in use.")
  tf.app.flags.DEFINE_integer("beam_size", 1, "Beam size during decoding. "
                              "If 0, no decoder, the non-extended Neural GPU.")
  tf.app.flags.DEFINE_integer("max_target_vocab", 0,
                              "Maximal size of target vocabulary.")
  tf.app.flags.DEFINE_integer("decode_offset", 0, "Offset for decoding.")
  tf.app.flags.DEFINE_integer("task", -1, "Task id when running on borg.")
  tf.app.flags.DEFINE_integer("nprint", 0, "How many test examples to print out.")
  tf.app.flags.DEFINE_integer("eval_bin_print", 3, "How many bins step in eval.")
  tf.app.flags.DEFINE_integer("mode", 0, "Mode: 0-train other-decode.")
  tf.app.flags.DEFINE_bool("atrous", False, "Whether to use atrous convs.")
  tf.app.flags.DEFINE_bool("layer_norm", False, "Do layer normalization.")
  tf.app.flags.DEFINE_bool("quantize", False, "Whether to quantize variables.")
  tf.app.flags.DEFINE_bool("do_train", True, "If false, only update memory.")
  tf.app.flags.DEFINE_bool("rnn_baseline", False, "If true build an RNN instead.")
  tf.app.flags.DEFINE_bool("simple_tokenizer", False,
                           "If true, tokenize on spaces only, digits are 0.")
  tf.app.flags.DEFINE_bool("normalize_digits", True,
                           "Whether to normalize digits with simple tokenizer.")
  tf.app.flags.DEFINE_integer("vocab_size", 16, "Joint vocabulary size.")
  tf.app.flags.DEFINE_string("data_dir", "/tmp", "Data directory")
  tf.app.flags.DEFINE_string("train_dir", "/tmp/", "Directory to store models.")
  tf.app.flags.DEFINE_string("test_file_prefix", "", "Files to test (.en,.fr).")
  tf.app.flags.DEFINE_integer("max_train_data_size", 0,
                              "Limit on the size of training data (0: no limit).")
  tf.app.flags.DEFINE_string("word_vector_file_en", "",
                             "Optional file with word vectors to start training.")
  tf.app.flags.DEFINE_string("word_vector_file_fr", "",
                             "Optional file with word vectors to start training.")
  tf.app.flags.DEFINE_string("problem", "wmt", "What problem are we solving?.")
  
  tf.app.flags.DEFINE_integer("ps_tasks", 0, "Number of ps tasks used.")
  tf.app.flags.DEFINE_string("master", "", "Name of the TensorFlow master.")
  
  FLAGS = tf.app.flags.FLAGS
  EXTRA_EVAL = 10
  EVAL_LEN_INCR = 8
  MAXLEN_F = 2.0
  
  
  def zero_split(tok_list, append=None):
    """Split tok_list (list of ints) on 0s, append int to all parts if given."""
    res, cur, l = [], [], 0
    for tok in tok_list:
      if tok == 0:
        if append is not None:
          cur.append(append)
        res.append(cur)
        l = max(l, len(cur))
        cur = []
      else:
        cur.append(tok)
    if append is not None:
      cur.append(append)
    res.append(cur)
    l = max(l, len(cur))
    return res, l
  
  
  def read_data(source_path, target_path, buckets, max_size=None, print_out=True):
    """Read data from source and target files and put into buckets.
  
    Args:
      source_path: path to the files with token-ids for the source language.
      target_path: path to the file with token-ids for the target language;
        it must be aligned with the source file: n-th line contains the desired
        output for n-th line from the source_path.
      buckets: the buckets to use.
      max_size: maximum number of lines to read, all other will be ignored;
        if 0 or None, data files will be read completely (no limit).
        If set to 1, no data will be returned (empty lists of the right form).
      print_out: whether to print out status or not.
  
    Returns:
      data_set: a list of length len(_buckets); data_set[n] contains a list of
        (source, target) pairs read from the provided data files that fit
        into the n-th bucket, i.e., such that len(source) < _buckets[n][0] and
        len(target) < _buckets[n][1]; source and target are lists of token-ids.
    """
    data_set = [[] for _ in buckets]
    counter = 0
    if max_size != 1:
      with tf.gfile.GFile(source_path, mode="r") as source_file:
        with tf.gfile.GFile(target_path, mode="r") as target_file:
          source, target = source_file.readline(), target_file.readline()
          while source and target and (not max_size or counter < max_size):
            counter += 1
            if counter % 100000 == 0 and print_out:
              print("  reading data line %d" % counter)
              sys.stdout.flush()
            source_ids = [int(x) for x in source.split()]
            target_ids = [int(x) for x in target.split()]
            source_ids, source_len = zero_split(source_ids)
            target_ids, target_len = zero_split(target_ids, append=wmt.EOS_ID)
            for bucket_id, size in enumerate(buckets):
              if source_len <= size and target_len <= size:
                data_set[bucket_id].append([source_ids, target_ids])
                break
            source, target = source_file.readline(), target_file.readline()
    return data_set
  
  
  global_train_set = {"wmt": []}
  train_buckets_scale = {"wmt": []}
  
  
  def calculate_buckets_scale(data_set, buckets, problem):
    """Calculate buckets scales for the given data set."""
    train_bucket_sizes = [len(data_set[b]) for b in xrange(len(buckets))]
    train_total_size = max(1, float(sum(train_bucket_sizes)))
  
    # A bucket scale is a list of increasing numbers from 0 to 1 that we'll use
    # to select a bucket. Length of [scale[i], scale[i+1]] is proportional to
    # the size if i-th training bucket, as used later.
    if problem not in train_buckets_scale:
      train_buckets_scale[problem] = []
    train_buckets_scale[problem].append(
        [sum(train_bucket_sizes[:i + 1]) / train_total_size
         for i in xrange(len(train_bucket_sizes))])
    return train_total_size
  
  
  def read_data_into_global(source_path, target_path, buckets,
                            max_size=None, print_out=True):
    """Read data into the global variables (can be in a separate thread)."""
    # pylint: disable=global-variable-not-assigned
    global global_train_set, train_buckets_scale
    # pylint: enable=global-variable-not-assigned
    data_set = read_data(source_path, target_path, buckets, max_size, print_out)
    global_train_set["wmt"].append(data_set)
    train_total_size = calculate_buckets_scale(data_set, buckets, "wmt")
    if print_out:
      print("  Finished global data reading (%d)." % train_total_size)
  
  
  def initialize(sess=None):
    """Initialize data and model."""
    global MAXLEN_F
    # Create training directory if it does not exist.
    if not tf.gfile.IsDirectory(FLAGS.train_dir):
      data.print_out("Creating training directory %s." % FLAGS.train_dir)
      tf.gfile.MkDir(FLAGS.train_dir)
    decode_suffix = "beam%dln%d" % (FLAGS.beam_size,
                                    int(100 * FLAGS.length_norm))
    if FLAGS.mode == 0:
      decode_suffix = ""
    if FLAGS.task >= 0:
      data.log_filename = os.path.join(FLAGS.train_dir,
                                       "log%d%s" % (FLAGS.task, decode_suffix))
    else:
      data.log_filename = os.path.join(FLAGS.train_dir, "neural_gpu/log")
  
    # Set random seed.
    if FLAGS.random_seed > 0:
      seed = FLAGS.random_seed + max(0, FLAGS.task)
      tf.set_random_seed(seed)
      random.seed(seed)
      np.random.seed(seed)
  
    # Check data sizes.
    assert data.bins
    max_length = min(FLAGS.max_length, data.bins[-1])
    while len(data.bins) > 1 and data.bins[-2] >= max_length + EXTRA_EVAL:
      data.bins = data.bins[:-1]
    if sess is None and FLAGS.task == 0 and FLAGS.num_replicas > 1:
      if max_length > 60:
        max_length = max_length * 1 / 2  # Save memory on chief.
    min_length = min(14, max_length - 3) if FLAGS.problem == "wmt" else 3
    for p in FLAGS.problem.split("-"):
      if p in ["progeval", "progsynth"]:
        min_length = max(26, min_length)
    assert max_length + 1 > min_length
    while len(data.bins) > 1 and data.bins[-2] >= max_length + EXTRA_EVAL:
      data.bins = data.bins[:-1]
  
    # Create checkpoint directory if it does not exist.
    if FLAGS.mode == 0 or FLAGS.task < 0:
      checkpoint_dir = os.path.join(FLAGS.train_dir, "neural_gpu%s"
                                    % ("" if FLAGS.task < 0 else str(FLAGS.task)))
    else:
      checkpoint_dir = FLAGS.train_dir
    if not tf.gfile.IsDirectory(checkpoint_dir):
      data.print_out("Creating checkpoint directory %s." % checkpoint_dir)
      tf.gfile.MkDir(checkpoint_dir)
  
    # Prepare data.
    if FLAGS.problem == "wmt":
      # Prepare WMT data.
      data.print_out("Preparing WMT data in %s" % FLAGS.data_dir)
      if FLAGS.simple_tokenizer:
        MAXLEN_F = 3.5
        (en_train, fr_train, en_dev, fr_dev,
         en_path, fr_path) = wmt.prepare_wmt_data(
             FLAGS.data_dir, FLAGS.vocab_size,
             tokenizer=wmt.space_tokenizer,
             normalize_digits=FLAGS.normalize_digits)
      else:
        (en_train, fr_train, en_dev, fr_dev,
         en_path, fr_path) = wmt.prepare_wmt_data(
             FLAGS.data_dir, FLAGS.vocab_size)
  
      # Read data into buckets and compute their sizes.
      fr_vocab, rev_fr_vocab = wmt.initialize_vocabulary(fr_path)
      data.vocab = fr_vocab
      data.rev_vocab = rev_fr_vocab
      data.print_out("Reading development and training data (limit: %d)."
                     % FLAGS.max_train_data_size)
      dev_set = {}
      dev_set["wmt"] = read_data(en_dev, fr_dev, data.bins)
      def data_read(size, print_out):
        read_data_into_global(en_train, fr_train, data.bins, size, print_out)
      data_read(50000, False)
      read_thread_small = threading.Thread(
          name="reading-data-small", target=lambda: data_read(900000, False))
      read_thread_small.start()
      read_thread_full = threading.Thread(
          name="reading-data-full",
          target=lambda: data_read(FLAGS.max_train_data_size, True))
      read_thread_full.start()
      data.print_out("Data reading set up.")
    else:
      # Prepare algorithmic data.
      en_path, fr_path = None, None
      tasks = FLAGS.problem.split("-")
      data_size = FLAGS.train_data_size
      for t in tasks:
        data.print_out("Generating data for %s." % t)
        if t in ["progeval", "progsynth"]:
          data.init_data(t, data.bins[-1], 20 * data_size, FLAGS.vocab_size)
          if len(program_utils.prog_vocab) > FLAGS.vocab_size - 2:
            raise ValueError("Increase vocab_size to %d for prog-tasks."
                             % (len(program_utils.prog_vocab) + 2))
          data.rev_vocab = program_utils.prog_vocab
          data.vocab = program_utils.prog_rev_vocab
        else:
          for l in xrange(max_length + EXTRA_EVAL - 1):
            data.init_data(t, l, data_size, FLAGS.vocab_size)
          data.init_data(t, data.bins[-2], data_size, FLAGS.vocab_size)
          data.init_data(t, data.bins[-1], data_size, FLAGS.vocab_size)
        if t not in global_train_set:
          global_train_set[t] = []
        global_train_set[t].append(data.train_set[t])
        calculate_buckets_scale(data.train_set[t], data.bins, t)
      dev_set = data.test_set
  
    # Grid-search parameters.
    lr = FLAGS.lr
    init_weight = FLAGS.init_weight
    max_grad_norm = FLAGS.max_grad_norm
    if sess is not None and FLAGS.task > -1:
      def job_id_factor(step):
        """If jobid / step mod 3 is 0, 1, 2: say 0, 1, -1."""
        return ((((FLAGS.task / step) % 3) + 1) % 3) - 1
      lr *= math.pow(2, job_id_factor(1))
      init_weight *= math.pow(1.5, job_id_factor(3))
      max_grad_norm *= math.pow(2, job_id_factor(9))
  
    # Print out parameters.
    curriculum = FLAGS.curriculum_seq
    msg1 = ("layers %d kw %d h %d kh %d batch %d noise %.2f"
            % (FLAGS.nconvs, FLAGS.kw, FLAGS.height, FLAGS.kh,
               FLAGS.batch_size, FLAGS.grad_noise_scale))
    msg2 = ("cut %.2f lr %.3f iw %.2f cr %.2f nm %d d%.4f gn %.2f %s"
            % (FLAGS.cutoff, lr, init_weight, curriculum, FLAGS.nmaps,
               FLAGS.dropout, max_grad_norm, msg1))
    data.print_out(msg2)
  
    # Create model and initialize it.
    tf.get_variable_scope().set_initializer(
        tf.orthogonal_initializer(gain=1.8 * init_weight))
    max_sampling_rate = FLAGS.max_sampling_rate if FLAGS.mode == 0 else 0.0
    o = FLAGS.vocab_size if FLAGS.max_target_vocab < 1 else FLAGS.max_target_vocab
    ngpu.CHOOSE_K = FLAGS.soft_mem_size
    do_beam_model = FLAGS.train_beam_freq > 0.0001 and FLAGS.beam_size > 1
    beam_size = FLAGS.beam_size if FLAGS.mode > 0 and not do_beam_model else 1
    beam_size = min(beam_size, FLAGS.beam_size)
    beam_model = None
    def make_ngpu(cur_beam_size, back):
      return ngpu.NeuralGPU(
          FLAGS.nmaps, FLAGS.vec_size, FLAGS.vocab_size, o,
          FLAGS.dropout, max_grad_norm, FLAGS.cutoff, FLAGS.nconvs,
          FLAGS.kw, FLAGS.kh, FLAGS.height, FLAGS.mem_size,
          lr / math.sqrt(FLAGS.num_replicas), min_length + 3, FLAGS.num_gpus,
          FLAGS.num_replicas, FLAGS.grad_noise_scale, max_sampling_rate,
          atrous=FLAGS.atrous, do_rnn=FLAGS.rnn_baseline,
          do_layer_norm=FLAGS.layer_norm, beam_size=cur_beam_size, backward=back)
    if sess is None:
      with tf.device(tf.train.replica_device_setter(FLAGS.ps_tasks)):
        model = make_ngpu(beam_size, True)
        if do_beam_model:
          tf.get_variable_scope().reuse_variables()
          beam_model = make_ngpu(FLAGS.beam_size, False)
    else:
      model = make_ngpu(beam_size, True)
      if do_beam_model:
        tf.get_variable_scope().reuse_variables()
        beam_model = make_ngpu(FLAGS.beam_size, False)
  
    sv = None
    if sess is None:
      # The supervisor configuration has a few overriden options.
      sv = tf.train.Supervisor(logdir=checkpoint_dir,
                               is_chief=(FLAGS.task < 1),
                               saver=model.saver,
                               summary_op=None,
                               save_summaries_secs=60,
                               save_model_secs=15 * 60,
                               global_step=model.global_step)
  
      config = tf.ConfigProto(allow_soft_placement=True)
      sess = sv.PrepareSession(FLAGS.master, config=config)
  
    data.print_out("Created model. Checkpoint dir %s" % checkpoint_dir)
  
    # Load model from parameters if a checkpoint exists.
    ckpt = tf.train.get_checkpoint_state(checkpoint_dir)
    if ckpt and tf.gfile.Exists(ckpt.model_checkpoint_path + ".index"):
      data.print_out("Reading model parameters from %s"
                     % ckpt.model_checkpoint_path)
      model.saver.restore(sess, ckpt.model_checkpoint_path)
    elif sv is None:
      sess.run(tf.global_variables_initializer())
      data.print_out("Initialized variables (no supervisor mode).")
    elif FLAGS.task < 1 and FLAGS.mem_size > 0:
      # sess.run(model.mem_norm_op)
      data.print_out("Created new model and normalized mem (on chief).")
  
    # Return the model and needed variables.
    return (model, beam_model, min_length, max_length, checkpoint_dir,
            (global_train_set, dev_set, en_path, fr_path), sv, sess)
  
  
  def m_step(model, beam_model, sess, batch_size, inp, target, bucket, nsteps, p):
    """Evaluation multi-step for program synthesis."""
    state, scores, hist = None, [[-11.0 for _ in xrange(batch_size)]], []
    for _ in xrange(nsteps):
      # Get the best beam (no training, just forward model).
      new_target, new_first, new_inp, new_scores = get_best_beam(
          beam_model, sess, inp, target,
          batch_size, FLAGS.beam_size, bucket, hist, p, test_mode=True)
      hist.append(new_first)
      _, _, _, state = model.step(sess, inp, new_target, False, state=state)
      inp = new_inp
      scores.append([max(scores[-1][i], new_scores[i])
                     for i in xrange(batch_size)])
    # The final step with the true target.
    loss, res, _, _ = model.step(sess, inp, target, False, state=state)
    return loss, res, new_target, scores[1:]
  
  
  def single_test(bin_id, model, sess, nprint, batch_size, dev, p, print_out=True,
                  offset=None, beam_model=None):
    """Test model on test data of length l using the given session."""
    if not dev[p][bin_id]:
      data.print_out("  bin %d (%d)\t%s\tppl NA errors NA seq-errors NA"
                     % (bin_id, data.bins[bin_id], p))
      return 1.0, 1.0, 0.0
    inpt, target = data.get_batch(
        bin_id, batch_size, dev[p], FLAGS.height, offset)
    if FLAGS.beam_size > 1 and beam_model:
      loss, res, new_tgt, scores = m_step(
          model, beam_model, sess, batch_size, inpt, target, bin_id,
          FLAGS.eval_beam_steps, p)
      score_avgs = [sum(s) / float(len(s)) for s in scores]
      score_maxs = [max(s) for s in scores]
      score_str = ["(%.2f, %.2f)" % (score_avgs[i], score_maxs[i])
                   for i in xrange(FLAGS.eval_beam_steps)]
      data.print_out("  == scores (avg, max): %s" % "; ".join(score_str))
      errors, total, seq_err = data.accuracy(inpt, res, target, batch_size,
                                             nprint, new_tgt, scores[-1])
    else:
      loss, res, _, _ = model.step(sess, inpt, target, False)
      errors, total, seq_err = data.accuracy(inpt, res, target, batch_size,
                                             nprint)
    seq_err = float(seq_err) / batch_size
    if total > 0:
      errors = float(errors) / total
    if print_out:
      data.print_out("  bin %d (%d)\t%s\tppl %.2f errors %.2f seq-errors %.2f"
                     % (bin_id, data.bins[bin_id], p, data.safe_exp(loss),
                        100 * errors, 100 * seq_err))
    return (errors, seq_err, loss)
  
  
  def assign_vectors(word_vector_file, embedding_key, vocab_path, sess):
    """Assign the embedding_key variable from the given word vectors file."""
    # For words in the word vector file, set their embedding at start.
    if not tf.gfile.Exists(word_vector_file):
      data.print_out("Word vector file does not exist: %s" % word_vector_file)
      sys.exit(1)
    vocab, _ = wmt.initialize_vocabulary(vocab_path)
    vectors_variable = [v for v in tf.trainable_variables()
                        if embedding_key == v.name]
    if len(vectors_variable) != 1:
      data.print_out("Word vector variable not found or too many.")
      sys.exit(1)
    vectors_variable = vectors_variable[0]
    vectors = vectors_variable.eval()
    data.print_out("Pre-setting word vectors from %s" % word_vector_file)
    with tf.gfile.GFile(word_vector_file, mode="r") as f:
      # Lines have format: dog 0.045123 -0.61323 0.413667 ...
      for line in f:
        line_parts = line.split()
        # The first part is the word.
        word = line_parts[0]
        if word in vocab:
          # Remaining parts are components of the vector.
          word_vector = np.array(map(float, line_parts[1:]))
          if len(word_vector) != FLAGS.vec_size:
            data.print_out("Warn: Word '%s', Expecting vector size %d, "
                           "found %d" % (word, FLAGS.vec_size,
                                         len(word_vector)))
          else:
            vectors[vocab[word]] = word_vector
    # Assign the modified vectors to the vectors_variable in the graph.
    sess.run([vectors_variable.initializer],
             {vectors_variable.initializer.inputs[1]: vectors})
  
  
  def print_vectors(embedding_key, vocab_path, word_vector_file):
    """Print vectors from the given variable."""
    _, rev_vocab = wmt.initialize_vocabulary(vocab_path)
    vectors_variable = [v for v in tf.trainable_variables()
                        if embedding_key == v.name]
    if len(vectors_variable) != 1:
      data.print_out("Word vector variable not found or too many.")
      sys.exit(1)
    vectors_variable = vectors_variable[0]
    vectors = vectors_variable.eval()
    l, s = vectors.shape[0], vectors.shape[1]
    data.print_out("Printing %d word vectors from %s to %s."
                   % (l, embedding_key, word_vector_file))
    with tf.gfile.GFile(word_vector_file, mode="w") as f:
      # Lines have format: dog 0.045123 -0.61323 0.413667 ...
      for i in xrange(l):
        f.write(rev_vocab[i])
        for j in xrange(s):
          f.write(" %.8f" % vectors[i][j])
        f.write("\n")
  
  
  def get_bucket_id(train_buckets_scale_c, max_cur_length, data_set):
    """Get a random bucket id."""
    # Choose a bucket according to data distribution. Pick a random number
    # in [0, 1] and use the corresponding interval in train_buckets_scale.
    random_number_01 = np.random.random_sample()
    bucket_id = min([i for i in xrange(len(train_buckets_scale_c))
                     if train_buckets_scale_c[i] > random_number_01])
    while bucket_id > 0 and not data_set[bucket_id]:
      bucket_id -= 1
    for _ in xrange(10 if np.random.random_sample() < 0.9 else 1):
      if data.bins[bucket_id] > max_cur_length:
        random_number_01 = min(random_number_01, np.random.random_sample())
        bucket_id = min([i for i in xrange(len(train_buckets_scale_c))
                         if train_buckets_scale_c[i] > random_number_01])
        while bucket_id > 0 and not data_set[bucket_id]:
          bucket_id -= 1
    return bucket_id
  
  
  def score_beams(beams, target, inp, history, p,
                  print_out=False, test_mode=False):
    """Score beams."""
    if p == "progsynth":
      return score_beams_prog(beams, target, inp, history, print_out, test_mode)
    elif test_mode:
      return beams[0], 10.0 if str(beams[0][:len(target)]) == str(target) else 0.0
    else:
      history_s = [str(h) for h in history]
      best, best_score, tgt, eos_id = None, -1000.0, target, None
      if p == "wmt":
        eos_id = wmt.EOS_ID
      if eos_id and eos_id in target:
        tgt = target[:target.index(eos_id)]
      for beam in beams:
        if eos_id and eos_id in beam:
          beam = beam[:beam.index(eos_id)]
        l = min(len(tgt), len(beam))
        score = len([i for i in xrange(l) if tgt[i] == beam[i]]) / float(len(tgt))
        hist_score = 20.0 if str([b for b in beam if b > 0]) in history_s else 0.0
        if score < 1.0:
          score -= hist_score
        if score > best_score:
          best = beam
          best_score = score
      return best, best_score
  
  
  def score_beams_prog(beams, target, inp, history, print_out=False,
                       test_mode=False):
    """Score beams for program synthesis."""
    tgt_prog = linearize(target, program_utils.prog_vocab, True, 1)
    hist_progs = [linearize(h, program_utils.prog_vocab, True, 1)
                  for h in history]
    tgt_set = set(target)
    if print_out:
      print("target: ", tgt_prog)
    inps, tgt_outs = [], []
    for i in xrange(3):
      ilist = [inp[i + 1, l] for l in xrange(inp.shape[1])]
      clist = [program_utils.prog_vocab[x] for x in ilist if x > 0]
      olist = clist[clist.index("]") + 1:]  # outputs
      clist = clist[1:clist.index("]")]     # inputs
      inps.append([int(x) for x in clist])
      if olist[0] == "[":  # olist may be [int] or just int
        tgt_outs.append(str([int(x) for x in olist[1:-1]]))
      else:
        if len(olist) == 1:
          tgt_outs.append(olist[0])
        else:
          print([program_utils.prog_vocab[x] for x in ilist if x > 0])
          print(olist)
          print(tgt_prog)
          print(program_utils.evaluate(tgt_prog, {"a": inps[-1]}))
          print("AAAAA")
          tgt_outs.append(olist[0])
    if not test_mode:
      for _ in xrange(7):
        ilen = np.random.randint(len(target) - 3) + 1
        inps.append([random.choice(range(-15, 15)) for _ in range(ilen)])
      tgt_outs.extend([program_utils.evaluate(tgt_prog, {"a": inp})
                       for inp in inps[3:]])
    best, best_prog, best_score = None, "", -1000.0
    for beam in beams:
      b_prog = linearize(beam, program_utils.prog_vocab, True, 1)
      b_set = set(beam)
      jsim = len(tgt_set & b_set) / float(len(tgt_set | b_set))
      b_outs = [program_utils.evaluate(b_prog, {"a": inp}) for inp in inps]
      errs = len([x for x in b_outs if x == "ERROR"])
      imatches = len([i for i in xrange(3) if b_outs[i] == tgt_outs[i]])
      perfect = 10.0 if imatches == 3 else 0.0
      hist_score = 20.0 if b_prog in hist_progs else 0.0
      if test_mode:
        score = perfect - errs
      else:
        matches = len([i for i in xrange(10) if b_outs[i] == tgt_outs[i]])
        score = perfect + matches + jsim - errs
      if score < 10.0:
        score -= hist_score
      # print b_prog
      # print "jsim: ", jsim, " errs: ", errs, " mtchs: ", matches, " s: ", score
      if score > best_score:
        best = beam
        best_prog = b_prog
        best_score = score
    if print_out:
      print("best score: ", best_score, " best prog: ", best_prog)
    return best, best_score
  
  
  def get_best_beam(beam_model, sess, inp, target, batch_size, beam_size,
                    bucket, history, p, test_mode=False):
    """Run beam_model, score beams, and return the best as target and in input."""
    _, output_logits, _, _ = beam_model.step(
        sess, inp, target, None, beam_size=FLAGS.beam_size)
    new_targets, new_firsts, scores, new_inp = [], [], [], np.copy(inp)
    for b in xrange(batch_size):
      outputs = []
      history_b = [[h[b, 0, l] for l in xrange(data.bins[bucket])]
                   for h in history]
      for beam_idx in xrange(beam_size):
        outputs.append([int(o[beam_idx * batch_size + b])
                        for o in output_logits])
      target_t = [target[b, 0, l] for l in xrange(data.bins[bucket])]
      best, best_score = score_beams(
          outputs, [t for t in target_t if t > 0], inp[b, :, :],
          [[t for t in h if t > 0] for h in history_b], p, test_mode=test_mode)
      scores.append(best_score)
      if 1 in best:  # Only until _EOS.
        best = best[:best.index(1) + 1]
      best += [0 for _ in xrange(len(target_t) - len(best))]
      new_targets.append([best])
      first, _ = score_beams(
          outputs, [t for t in target_t if t > 0], inp[b, :, :],
          [[t for t in h if t > 0] for h in history_b], p, test_mode=True)
      if 1 in first:  # Only until _EOS.
        first = first[:first.index(1) + 1]
      first += [0 for _ in xrange(len(target_t) - len(first))]
      new_inp[b, 0, :] = np.array(first, dtype=np.int32)
      new_firsts.append([first])
    # Change target if we found a great answer.
    new_target = np.array(new_targets, dtype=np.int32)
    for b in xrange(batch_size):
      if scores[b] >= 10.0:
        target[b, 0, :] = new_target[b, 0, :]
    new_first = np.array(new_firsts, dtype=np.int32)
    return new_target, new_first, new_inp, scores
  
  
  def train():
    """Train the model."""
    batch_size = FLAGS.batch_size * FLAGS.num_gpus
    (model, beam_model, min_length, max_length, checkpoint_dir,
     (train_set, dev_set, en_vocab_path, fr_vocab_path), sv, sess) = initialize()
    with sess.as_default():
      quant_op = model.quantize_op
      max_cur_length = min(min_length + 3, max_length)
      prev_acc_perp = [1000000 for _ in xrange(5)]
      prev_seq_err = 1.0
      is_chief = FLAGS.task < 1
      do_report = False
  
      # Main traning loop.
      while not sv.ShouldStop():
        global_step, max_cur_length, learning_rate = sess.run(
            [model.global_step, model.cur_length, model.lr])
        acc_loss, acc_l1, acc_total, acc_errors, acc_seq_err = 0.0, 0.0, 0, 0, 0
        acc_grad_norm, step_count, step_c1, step_time = 0.0, 0, 0, 0.0
  
        # For words in the word vector file, set their embedding at start.
        bound1 = FLAGS.steps_per_checkpoint - 1
        if FLAGS.word_vector_file_en and global_step < bound1 and is_chief:
          assign_vectors(FLAGS.word_vector_file_en, "embedding:0",
                         en_vocab_path, sess)
          if FLAGS.max_target_vocab < 1:
            assign_vectors(FLAGS.word_vector_file_en, "target_embedding:0",
                           en_vocab_path, sess)
  
        if FLAGS.word_vector_file_fr and global_step < bound1 and is_chief:
          assign_vectors(FLAGS.word_vector_file_fr, "embedding:0",
                         fr_vocab_path, sess)
          if FLAGS.max_target_vocab < 1:
            assign_vectors(FLAGS.word_vector_file_fr, "target_embedding:0",
                           fr_vocab_path, sess)
  
        for _ in xrange(FLAGS.steps_per_checkpoint):
          step_count += 1
          step_c1 += 1
          global_step = int(model.global_step.eval())
          train_beam_anneal = global_step / float(FLAGS.train_beam_anneal)
          train_beam_freq = FLAGS.train_beam_freq * min(1.0, train_beam_anneal)
          p = random.choice(FLAGS.problem.split("-"))
          train_set = global_train_set[p][-1]
          bucket_id = get_bucket_id(train_buckets_scale[p][-1], max_cur_length,
                                    train_set)
          # Prefer longer stuff 60% of time if not wmt.
          if np.random.randint(100) < 60 and FLAGS.problem != "wmt":
            bucket1 = get_bucket_id(train_buckets_scale[p][-1], max_cur_length,
                                    train_set)
            bucket_id = max(bucket1, bucket_id)
  
          # Run a step and time it.
          start_time = time.time()
          inp, target = data.get_batch(bucket_id, batch_size, train_set,
                                       FLAGS.height)
          noise_param = math.sqrt(math.pow(global_step + 1, -0.55) *
                                  prev_seq_err) * FLAGS.grad_noise_scale
          # In multi-step mode, we use best from beam for middle steps.
          state, new_target, scores, history = None, None, None, []
          while (FLAGS.beam_size > 1 and
                 train_beam_freq > np.random.random_sample()):
            # Get the best beam (no training, just forward model).
            new_target, new_first, new_inp, scores = get_best_beam(
                beam_model, sess, inp, target,
                batch_size, FLAGS.beam_size, bucket_id, history, p)
            history.append(new_first)
            # Training step with the previous input and the best beam as target.
            _, _, _, state = model.step(sess, inp, new_target, FLAGS.do_train,
                                        noise_param, update_mem=True, state=state)
            # Change input to the new one for the next step.
            inp = new_inp
            # If all results are great, stop (todo: not to wait for all?).
            if FLAGS.nprint > 1:
              print(scores)
            if sum(scores) / float(len(scores)) >= 10.0:
              break
          # The final step with the true target.
          loss, res, gnorm, _ = model.step(
              sess, inp, target, FLAGS.do_train, noise_param,
              update_mem=True, state=state)
          step_time += time.time() - start_time
          acc_grad_norm += 0.0 if gnorm is None else float(gnorm)
  
          # Accumulate statistics.
          acc_loss += loss
          acc_l1 += loss
          errors, total, seq_err = data.accuracy(
              inp, res, target, batch_size, 0, new_target, scores)
          if FLAGS.nprint > 1:
            print("seq_err: ", seq_err)
          acc_total += total
          acc_errors += errors
          acc_seq_err += seq_err
  
          # Report summary every 10 steps.
          if step_count + 3 > FLAGS.steps_per_checkpoint:
            do_report = True  # Don't polute plot too early.
          if is_chief and step_count % 10 == 1 and do_report:
            cur_loss = acc_l1 / float(step_c1)
            acc_l1, step_c1 = 0.0, 0
            cur_perp = data.safe_exp(cur_loss)
            summary = tf.Summary()
            summary.value.extend(
                [tf.Summary.Value(tag="log_perplexity", simple_value=cur_loss),
                 tf.Summary.Value(tag="perplexity", simple_value=cur_perp)])
            sv.SummaryComputed(sess, summary, global_step)
  
        # Normalize and print out accumulated statistics.
        acc_loss /= step_count
        step_time /= FLAGS.steps_per_checkpoint
        acc_seq_err = float(acc_seq_err) / (step_count * batch_size)
        prev_seq_err = max(0.0, acc_seq_err - 0.02)  # No noise at error < 2%.
        acc_errors = float(acc_errors) / acc_total if acc_total > 0 else 1.0
        t_size = float(sum([len(x) for x in train_set])) / float(1000000)
        msg = ("step %d step-time %.2f train-size %.3f lr %.6f grad-norm %.4f"
               % (global_step + 1, step_time, t_size, learning_rate,
                  acc_grad_norm / FLAGS.steps_per_checkpoint))
        data.print_out("%s len %d ppl %.6f errors %.2f sequence-errors %.2f" %
                       (msg, max_cur_length, data.safe_exp(acc_loss),
                        100*acc_errors, 100*acc_seq_err))
  
        # If errors are below the curriculum threshold, move curriculum forward.
        is_good = FLAGS.curriculum_ppx > data.safe_exp(acc_loss)
        is_good = is_good and FLAGS.curriculum_seq > acc_seq_err
        if is_good and is_chief:
          if FLAGS.quantize:
            # Quantize weights.
            data.print_out("  Quantizing parameters.")
            sess.run([quant_op])
          # Increase current length (until the next with training data).
          sess.run(model.cur_length_incr_op)
          # Forget last perplexities if we're not yet at the end.
          if max_cur_length < max_length:
            prev_acc_perp.append(1000000)
  
        # Lower learning rate if we're worse than the last 5 checkpoints.
        acc_perp = data.safe_exp(acc_loss)
        if acc_perp > max(prev_acc_perp[-5:]) and is_chief:
          sess.run(model.lr_decay_op)
        prev_acc_perp.append(acc_perp)
  
        # Save checkpoint.
        if is_chief:
          checkpoint_path = os.path.join(checkpoint_dir, "neural_gpu.ckpt")
          model.saver.save(sess, checkpoint_path,
                           global_step=model.global_step)
  
          # Run evaluation.
          bin_bound = 4
          for p in FLAGS.problem.split("-"):
            total_loss, total_err, tl_counter = 0.0, 0.0, 0
            for bin_id in xrange(len(data.bins)):
              if bin_id < bin_bound or bin_id % FLAGS.eval_bin_print == 1:
                err, _, loss = single_test(bin_id, model, sess, FLAGS.nprint,
                                           batch_size * 4, dev_set, p,
                                           beam_model=beam_model)
                if loss > 0.0:
                  total_loss += loss
                  total_err += err
                  tl_counter += 1
            test_loss = total_loss / max(1, tl_counter)
            test_err = total_err / max(1, tl_counter)
            test_perp = data.safe_exp(test_loss)
            summary = tf.Summary()
            summary.value.extend(
                [tf.Summary.Value(tag="test/%s/loss" % p, simple_value=test_loss),
                 tf.Summary.Value(tag="test/%s/error" % p, simple_value=test_err),
                 tf.Summary.Value(tag="test/%s/perplexity" % p,
                                  simple_value=test_perp)])
            sv.SummaryComputed(sess, summary, global_step)
  
  
  def linearize(output, rev_fr_vocab, simple_tokenizer=None, eos_id=wmt.EOS_ID):
    # If there is an EOS symbol in outputs, cut them at that point (WMT).
    if eos_id in output:
      output = output[:output.index(eos_id)]
    # Print out French sentence corresponding to outputs.
    if simple_tokenizer or FLAGS.simple_tokenizer:
      vlen = len(rev_fr_vocab)
      def vget(o):
        if o < vlen:
          return rev_fr_vocab[o]
        return "UNK"
      return " ".join([vget(o) for o in output])
    else:
      return wmt.basic_detokenizer([rev_fr_vocab[o] for o in output])
  
  
  def evaluate():
    """Evaluate an existing model."""
    batch_size = FLAGS.batch_size * FLAGS.num_gpus
    with tf.Session(config=tf.ConfigProto(allow_soft_placement=True)) as sess:
      (model, beam_model, _, _, _,
       (_, dev_set, en_vocab_path, fr_vocab_path), _, sess) = initialize(sess)
      for p in FLAGS.problem.split("-"):
        for bin_id in xrange(len(data.bins)):
          if (FLAGS.task >= 0 and bin_id > 4) or (FLAGS.nprint == 0 and
                                                  bin_id > 8 and p == "wmt"):
            break
          single_test(bin_id, model, sess, FLAGS.nprint, batch_size, dev_set, p,
                      beam_model=beam_model)
      path = FLAGS.test_file_prefix
      xid = "" if FLAGS.task < 0 else ("%.4d" % (FLAGS.task+FLAGS.decode_offset))
      en_path, fr_path = path + ".en" + xid, path + ".fr" + xid
      # Evaluate the test file if they exist.
      if path and tf.gfile.Exists(en_path) and tf.gfile.Exists(fr_path):
        data.print_out("Translating test set %s" % en_path)
        # Read lines.
        en_lines, fr_lines = [], []
        with tf.gfile.GFile(en_path, mode="r") as f:
          for line in f:
            en_lines.append(line.strip())
        with tf.gfile.GFile(fr_path, mode="r") as f:
          for line in f:
            fr_lines.append(line.strip())
        # Tokenize and convert to ids.
        en_vocab, _ = wmt.initialize_vocabulary(en_vocab_path)
        _, rev_fr_vocab = wmt.initialize_vocabulary(fr_vocab_path)
        if FLAGS.simple_tokenizer:
          en_ids = [wmt.sentence_to_token_ids(
              l, en_vocab, tokenizer=wmt.space_tokenizer,
              normalize_digits=FLAGS.normalize_digits)
                    for l in en_lines]
        else:
          en_ids = [wmt.sentence_to_token_ids(l, en_vocab) for l in en_lines]
        # Translate.
        results = []
        for idx, token_ids in enumerate(en_ids):
          if idx % 5 == 0:
            data.print_out("Translating example %d of %d." % (idx, len(en_ids)))
          # Which bucket does it belong to?
          buckets = [b for b in xrange(len(data.bins))
                     if data.bins[b] >= len(token_ids)]
          if buckets:
            result, result_cost = [], 100000000.0
            for bucket_id in buckets:
              if data.bins[bucket_id] > MAXLEN_F * len(token_ids) + EVAL_LEN_INCR:
                break
              # Get a 1-element batch to feed the sentence to the model.
              used_batch_size = 1  # batch_size
              inp, target = data.get_batch(
                  bucket_id, used_batch_size, None, FLAGS.height,
                  preset=([token_ids], [[]]))
              loss, output_logits, _, _ = model.step(
                  sess, inp, target, None, beam_size=FLAGS.beam_size)
              outputs = [int(o[0]) for o in output_logits]
              loss = loss[0] - (data.bins[bucket_id] * FLAGS.length_norm)
              if FLAGS.simple_tokenizer:
                cur_out = outputs
                if wmt.EOS_ID in cur_out:
                  cur_out = cur_out[:cur_out.index(wmt.EOS_ID)]
                res_tags = [rev_fr_vocab[o] for o in cur_out]
                bad_words, bad_brack = wmt.parse_constraints(token_ids, res_tags)
                loss += 1000.0 * bad_words + 100.0 * bad_brack
              # print (bucket_id, loss)
              if loss < result_cost:
                result = outputs
                result_cost = loss
            final = linearize(result, rev_fr_vocab)
            results.append("%s\t%s\n" % (final, fr_lines[idx]))
            # print result_cost
            sys.stderr.write(results[-1])
            sys.stderr.flush()
          else:
            sys.stderr.write("TOOO_LONG\t%s\n" % fr_lines[idx])
            sys.stderr.flush()
        if xid:
          decode_suffix = "beam%dln%dn" % (FLAGS.beam_size,
                                           int(100 * FLAGS.length_norm))
          with tf.gfile.GFile(path + ".res" + decode_suffix + xid, mode="w") as f:
            for line in results:
              f.write(line)
  
  
  def mul(l):
    res = 1.0
    for s in l:
      res *= s
    return res
  
  
  def interactive():
    """Interactively probe an existing model."""
    with tf.Session(config=tf.ConfigProto(allow_soft_placement=True)) as sess:
      # Initialize model.
      (model, _, _, _, _, (_, _, en_path, fr_path), _, _) = initialize(sess)
      # Load vocabularies.
      en_vocab, rev_en_vocab = wmt.initialize_vocabulary(en_path)
      _, rev_fr_vocab = wmt.initialize_vocabulary(fr_path)
      # Print out vectors and variables.
      if FLAGS.nprint > 0 and FLAGS.word_vector_file_en:
        print_vectors("embedding:0", en_path, FLAGS.word_vector_file_en)
      if FLAGS.nprint > 0 and FLAGS.word_vector_file_fr:
        print_vectors("target_embedding:0", fr_path, FLAGS.word_vector_file_fr)
      total = 0
      for v in tf.trainable_variables():
        shape = v.get_shape().as_list()
        total += mul(shape)
        print(v.name, shape, mul(shape))
      print(total)
      # Start interactive loop.
      sys.stdout.write("Input to Neural GPU Translation Model.\n")
      sys.stdout.write("> ")
      sys.stdout.flush()
      inpt = sys.stdin.readline(), ""
      while inpt:
        cures = []
        # Get token-ids for the input sentence.
        if FLAGS.simple_tokenizer:
          token_ids = wmt.sentence_to_token_ids(
              inpt, en_vocab, tokenizer=wmt.space_tokenizer,
              normalize_digits=FLAGS.normalize_digits)
        else:
          token_ids = wmt.sentence_to_token_ids(inpt, en_vocab)
        print([rev_en_vocab[t] for t in token_ids])
        # Which bucket does it belong to?
        buckets = [b for b in xrange(len(data.bins))
                   if data.bins[b] >= max(len(token_ids), len(cures))]
        if cures:
          buckets = [buckets[0]]
        if buckets:
          result, result_cost = [], 10000000.0
          for bucket_id in buckets:
            if data.bins[bucket_id] > MAXLEN_F * len(token_ids) + EVAL_LEN_INCR:
              break
            glen = 1
            for gen_idx in xrange(glen):
              # Get a 1-element batch to feed the sentence to the model.
              inp, target = data.get_batch(
                  bucket_id, 1, None, FLAGS.height, preset=([token_ids], [cures]))
              loss, output_logits, _, _ = model.step(
                  sess, inp, target, None, beam_size=FLAGS.beam_size,
                  update_mem=False)
              # If it is a greedy decoder, outputs are argmaxes of output_logits.
              if FLAGS.beam_size > 1:
                outputs = [int(o) for o in output_logits]
              else:
                loss = loss[0] - (data.bins[bucket_id] * FLAGS.length_norm)
                outputs = [int(np.argmax(logit, axis=1))
                           for logit in output_logits]
              print([rev_fr_vocab[t] for t in outputs])
              print(loss, data.bins[bucket_id])
              print(linearize(outputs, rev_fr_vocab))
              cures.append(outputs[gen_idx])
              print(cures)
              print(linearize(cures, rev_fr_vocab))
            if FLAGS.simple_tokenizer:
              cur_out = outputs
              if wmt.EOS_ID in cur_out:
                cur_out = cur_out[:cur_out.index(wmt.EOS_ID)]
              res_tags = [rev_fr_vocab[o] for o in cur_out]
              bad_words, bad_brack = wmt.parse_constraints(token_ids, res_tags)
              loss += 1000.0 * bad_words + 100.0 * bad_brack
            if loss < result_cost:
              result = outputs
              result_cost = loss
          print("FINAL", result_cost)
          print([rev_fr_vocab[t] for t in result])
          print(linearize(result, rev_fr_vocab))
        else:
          print("TOOO_LONG")
        sys.stdout.write("> ")
        sys.stdout.flush()
        inpt = sys.stdin.readline(), ""
  
  
  def main(_):
    if FLAGS.mode == 0:
      train()
    elif FLAGS.mode == 1:
      evaluate()
    else:
      interactive()
  
  if __name__ == "__main__":
    tf.app.run()
  