text: # -*- coding:utf-8 -*-
  
  import math
  import os
  import sys
  import time
  
  import numpy as np
  from six.moves import xrange
  import tensorflow as tf
  
  from nlp.textsum import model as seq2seq_model
  from nlp.textsum.dataset import data_utils, dataset
  
  
  def read_data(buckets, source_path, target_path):
  
      data_set = [[] for _ in buckets]
      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()
              counter = 0
              while source and target:
                  counter += 1
                  if counter % 10000 == 0:
                      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()]
                  target_ids.append(data_utils.EOS_ID)
                  for bucket_id, (source_size, target_size) in enumerate(buckets):
                      if len(source_ids) < source_size and len(target_ids) < target_size:
                          data_set[bucket_id].append([source_ids, target_ids])
                          break
                  source, target = source_file.readline(), target_file.readline()
      return data_set
  
  
  def train(args, steps_per_checkpoint=10):
      if not os.path.isdir(args.utils_dir):
          os.mkdir(args.utils_dir)
  
      if not os.path.isdir(args.train_dir):
          os.mkdir(args.train_dir)
  
      # Prepare Headline data.
      print("Preparing Headline data in %s" % args.data_dir)
      src_train, dest_train, src_dev, dest_dev, vocab_size =\
          data_utils.prepare_data(args.data_dir, args.utils_dir)
  
      args.vocab_size = vocab_size
  
      # device config for CPU usage
      config = tf.ConfigProto(device_count={"CPU": 4},  # limit to 4 CPU usage
                              inter_op_parallelism_threads=1,
                              intra_op_parallelism_threads=2)  # n threads parallel for ops
  
      with tf.Session(config=config) as sess:
          # Create model.
          print("Creating %d layers of %d units." % (args.num_layers, args.hidden_size))
          model = seq2seq_model.create_model(sess, args.train_dir, args, False)
  
          # Read data into buckets and compute their sizes.
          dev_set = read_data(args.buckets, src_dev, dest_dev)
          train_set = read_data(args.buckets, src_train, dest_train)
  
  
          train_bucket_sizes = [len(train_set[b]) for b in xrange(len(args.buckets))]
          train_total_size = float(sum(train_bucket_sizes))
  
          trainbuckets_scale = [sum(train_bucket_sizes[:i + 1]) / train_total_size
                                for i in xrange(len(train_bucket_sizes))]
  
          # This is the training loop.
          step_time, loss = 0.0, 0.0
          current_step = 0
          previous_losses = []
          while current_step<args.num_steps:
  
              random_number_01 = np.random.random_sample()
              bucket_id = min([i for i in xrange(len(trainbuckets_scale))
                               if trainbuckets_scale[i] > random_number_01])
  
              # Get a batch and make a step.
              start_time = time.time()
              encoder_inputs, decoder_inputs, target_weights =\
                  dataset.DataSet(args, train_set).next_batch(bucket_id)
  
              _, step_loss, _ = model.step(sess, encoder_inputs, decoder_inputs,
                                           target_weights, bucket_id, False)
              step_time += (time.time() - start_time) / steps_per_checkpoint
              loss += step_loss / steps_per_checkpoint
              current_step += 1
  
              # Once in a while, we save checkpoint, print statistics, and run evals.
              if current_step % steps_per_checkpoint == 0:
                  # Print statistics for the previous epoch.
                  perplexity = math.exp(float(loss)) if loss < 300 else float("inf")
                  print("global step %d learning rate %.4f step-time %.2f perplexity "
                        "%.2f" % (model.global_step.eval(), model.learning_rate.eval(),
                                  step_time, perplexity))
                  # Decrease learning rate if no improvement was seen over last 3 times.
                  if len(previous_losses) > 2 and loss > max(previous_losses[-3:]):
                      sess.run(model.learning_rate_decay_op)
                  previous_losses.append(loss)
                  # Save checkpoint and zero timer and loss.
                  checkpoint_path = os.path.join(args.train_dir, "headline_large.ckpt")
                  model.saver.save(sess, checkpoint_path, global_step=model.global_step)
                  step_time, loss = 0.0, 0.0
                  # Run evals on development set and print their perplexity.
                  for bucket_id in xrange(len(args.buckets)):
                      if len(dev_set[bucket_id]) == 0:
                          print("eval: empty bucket %d" % (bucket_id))
                          continue
                      encoder_inputs, decoder_inputs, target_weights =\
                          dataset.DataSet(args, dev_set).next_batch(bucket_id)
                      _, eval_loss, _ = model.step(sess, encoder_inputs, decoder_inputs,
                                                   target_weights, bucket_id, True)
                      eval_ppx = math.exp(float(eval_loss)) if eval_loss < 300 else float(
                          "inf")
                      print("  eval: bucket %d perplexity %.2f" % (bucket_id, eval_ppx))
                  sys.stdout.flush()