text: import os
  import json
  import math
  import random
  import sys
  import time
  
  import gflags
  import numpy as np
  
  from spinn.util import afs_safe_logger
  from spinn.util.data import SimpleProgressBar
  from spinn.util.blocks import to_gpu
  from spinn.util.misc import Accumulator, EvalReporter
  from spinn.util.logging import stats, train_accumulate, create_log_formatter
  from spinn.util.logging import train_rl_accumulate
  from spinn.util.logging import eval_stats, eval_accumulate, prettyprint_trees
  from spinn.util.loss import auxiliary_loss
  from spinn.util.sparks import sparks, dec_str
  import spinn.util.evalb as evalb
  import spinn.util.logging_pb2 as pb
  from spinn.util.trainer import ModelTrainer
  
  # PyTorch
  import torch
  import torch.nn as nn
  from torch.autograd import Variable
  
  from spinn.models.base import get_data_manager, get_flags, get_batch
  from spinn.models.base import flag_defaults, init_model, log_path
  from spinn.models.base import load_data_and_embeddings
  
  
  FLAGS = gflags.FLAGS
  
  
  def evaluate(FLAGS, model, eval_set, log_entry,
               logger, trainer, vocabulary=None, show_sample=False, eval_index=0):
      filename, dataset = eval_set
  
      A = Accumulator()
      eval_log = log_entry.evaluation.add()
      reporter = EvalReporter()
      tree_strs = None
  
      # Evaluate
      total_batches = len(dataset)
      progress_bar = SimpleProgressBar(
          msg="Run Eval",
          bar_length=60,
          enabled=FLAGS.show_progress_bar)
      progress_bar.step(0, total=total_batches)
      total_tokens = 0
      cpt = 0
      cpt_max = 0
      start = time.time()
  
      model.eval()
      for i, dataset_batch in enumerate(dataset):
          batch = get_batch(dataset_batch)
          eval_X_batch, eval_transitions_batch, eval_y_batch, eval_num_transitions_batch, eval_ids = batch
  
          # Run model.
          np.set_printoptions(threshold=np.inf)
          output = model(eval_X_batch, eval_transitions_batch, eval_y_batch,
                         use_internal_parser=FLAGS.use_internal_parser,
                         validate_transitions=FLAGS.validate_transitions,
                         store_parse_masks=show_sample,
                         example_lengths=eval_num_transitions_batch)
  
          can_sample = (FLAGS.model_type ==
                        "RLSPINN" and FLAGS.use_internal_parser)
          if show_sample and can_sample:
              tmp_samples = model.get_samples(
                  eval_X_batch, vocabulary, only_one=not FLAGS.write_eval_report)
              tree_strs = prettyprint_trees(tmp_samples)
          if not FLAGS.write_eval_report:
              # Only show one sample, regardless of the number of batches.
              show_sample = False
  
          # Calculate class accuracy.
          target = torch.from_numpy(eval_y_batch).long()
  
          # get the index of the max log-probability
          pred = output.data.max(1, keepdim=False)[1].cpu()
  
          eval_accumulate(model, A, batch)
          A.add('class_correct', pred.eq(target).sum())
          A.add('class_total', target.size(0))
  
          # Update Aggregate Accuracies
          total_tokens += sum(
              [(nt + 1) / 2 for nt in eval_num_transitions_batch.reshape(-1)])
  
          if FLAGS.write_eval_report:
              transitions_per_example, _ = model.spinn.get_transitions_per_example(
                  style="preds" if FLAGS.eval_report_use_preds else "given") if (
                  FLAGS.model_type == "RLSPINN" and FLAGS.use_internal_parser) else (
                  None, None)
  
              if model.use_sentence_pair:
                  batch_size = pred.size(0)
                  sent1_transitions = transitions_per_example[:
                                                              batch_size] if transitions_per_example is not None else None
                  sent2_transitions = transitions_per_example[batch_size:
                                                              ] if transitions_per_example is not None else None
  
                  sent1_trees = tree_strs[:batch_size] if tree_strs is not None else None
                  sent2_trees = tree_strs[batch_size:
                                          ] if tree_strs is not None else None
              else:
                  sent1_transitions = transitions_per_example if transitions_per_example is not None else None
                  sent2_transitions = None
  
                  sent1_trees = tree_strs if tree_strs is not None else None
                  sent2_trees = None
  
              if FLAGS.cp_metric:
                  cp, cp_max = reporter.save_batch(
                      pred,
                      target,
                      eval_ids,
                      output.data.cpu().numpy(),
                      sent1_transitions,
                      sent2_transitions,
                      sent1_trees,
                      sent2_trees, 
                      cp_metric=FLAGS.cp_metric,
                      mt=False)
                  cpt += cp
                  cpt_max += cp_max
              else:
                  reporter.save_batch(
                  pred,
                  target,
                  eval_ids,
                  output.data.cpu().numpy(),
                  sent1_transitions,
                  sent2_transitions,
                  sent1_trees,
                  sent2_trees,
                  mt=False)
  
          # Print Progress
          progress_bar.step(i + 1, total=total_batches)
      progress_bar.finish()
      
      cp_metric_value = cpt / cpt_max
  
      if tree_strs is not None:
          logger.Log('Sample: ' + tree_strs[0])
  
      end = time.time()
      total_time = end - start
  
      A.add('total_tokens', total_tokens)
      A.add('total_time', total_time)
  
      logger.Log("Eval cp_acc: " + str(cp_metric_value))
  
      eval_stats(model, A, eval_log)
      eval_log.filename = filename
  
      if FLAGS.write_eval_report:
          eval_report_path = os.path.join(
              FLAGS.log_path,
              FLAGS.experiment_name +
              ".eval_set_" +
              str(eval_index) +
              ".report")
          reporter.write_report(eval_report_path)
  
      eval_class_acc = eval_log.eval_class_accuracy
      eval_trans_acc = eval_log.eval_transition_accuracy
  
      return eval_class_acc, eval_trans_acc
  
  
  def train_loop(
          FLAGS,
          model,
          trainer,
          training_data_iter,
          eval_iterators,
          logger,
  	vocabulary):
      # Accumulate useful statistics.
      A = Accumulator(maxlen=FLAGS.deque_length)
  
      # Train.
      logger.Log("Training.")
      # New Training Loop
      progress_bar = SimpleProgressBar(
          msg="Training", bar_length=60, enabled=FLAGS.show_progress_bar)
      progress_bar.step(i=0, total=FLAGS.statistics_interval_steps)
  
      log_entry = pb.SpinnEntry()
      for _ in range(trainer.step, FLAGS.training_steps):
          if (trainer.step - trainer.best_dev_step) > FLAGS.early_stopping_steps_to_wait:
              logger.Log('No improvement after ' +
                         str(FLAGS.early_stopping_steps_to_wait) +
                         ' steps. Stopping training.')
              break
  
          model.train()
          log_entry.Clear()
          log_entry.step = trainer.step
          should_log = False
  
          start = time.time()
  
          batch = get_batch(next(training_data_iter))
          X_batch, transitions_batch, y_batch, num_transitions_batch, train_ids = batch
  
          total_tokens = sum(
              [(nt + 1) / 2 for nt in num_transitions_batch.reshape(-1)])
  
          # Reset cached gradients.
          trainer.optimizer_zero_grad()
  
          temperature = math.sin(
              math.pi /
              2 +
              trainer.step /
              float(
                  FLAGS.rl_confidence_interval) *
              2 *
              math.pi)
          temperature = (temperature + 1) / 2
  
          # Confidence Penalty for Transition Predictions.
          if FLAGS.rl_confidence_penalty:
              epsilon = FLAGS.rl_epsilon * \
                  math.exp(-trainer.step / float(FLAGS.rl_epsilon_decay))
              temp = 1 + \
                  (temperature - .5) * FLAGS.rl_confidence_penalty * epsilon
              model.spinn.temperature = max(1e-3, temp)
  
          # Soft Wake/Sleep based on temperature.
          if FLAGS.rl_wake_sleep:
              model.rl_weight = temperature * FLAGS.rl_weight
  
          # Run model.
          output = model(X_batch, transitions_batch, y_batch,
                         use_internal_parser=FLAGS.use_internal_parser,
                         validate_transitions=FLAGS.validate_transitions
                         )
          
          # Calculate class accuracy.
          target = torch.from_numpy(y_batch).long()
  
          # get the index of the max log-probability
          pred = output.data.max(1, keepdim=False)[1].cpu()
  
          class_acc = pred.eq(target).sum().float() / float(target.size(0))
  
          # Calculate class loss.
          xent_loss = nn.CrossEntropyLoss()(output, to_gpu(Variable(target)))
  
          # Optionally calculate transition loss.
          transition_loss = model.transition_loss if hasattr(
              model, 'transition_loss') else None
  
          # Accumulate Total Loss Variable
          total_loss = 0.0
          total_loss += xent_loss
          if transition_loss is not None and model.optimize_transition_loss:
              total_loss += transition_loss
          aux_loss = auxiliary_loss(model)
          total_loss += aux_loss[0]
  
          # Backward pass.
          total_loss.backward()
  
          # Hard Gradient Clipping
          nn.utils.clip_grad_norm_([param for name, param in model.named_parameters() if name not in ["embed.embed.weight"]], FLAGS.clipping_max_value)
  
          # Gradient descent step.
          trainer.optimizer_step()
  
          end = time.time()
  
          total_time = end - start
  
          train_accumulate(model, A, batch)
          A.add('class_acc', class_acc)
          A.add('total_tokens', total_tokens)
          A.add('total_time', total_time)
  
          train_rl_accumulate(model, A, batch)
  
          if trainer.step % FLAGS.statistics_interval_steps == 0:
              progress_bar.step(i=FLAGS.statistics_interval_steps,
                                total=FLAGS.statistics_interval_steps)
              progress_bar.finish()
  
              A.add('xent_cost', xent_loss.data.item())
              stats(model, trainer, A, log_entry)
              should_log = True
  
          if trainer.step % FLAGS.sample_interval_steps == 0 and FLAGS.num_samples > 0:
              should_log = True
              model.train()
  
              model(X_batch, transitions_batch, y_batch,
                use_internal_parser=FLAGS.use_internal_parser,
                validate_transitions=FLAGS.validate_transitions
                )
  
              tr_transitions_per_example, tr_strength = model.spinn.get_transitions_per_example(
              )
  
              model.eval()
              model(X_batch, transitions_batch, y_batch,
                    use_internal_parser=FLAGS.use_internal_parser,
                    validate_transitions=FLAGS.validate_transitions
                    )
              ev_transitions_per_example, ev_strength = model.spinn.get_transitions_per_example(
              )
  
              if model.use_sentence_pair and len(transitions_batch.shape) == 3:
                  transitions_batch = np.concatenate([
                      transitions_batch[:, :, 0], transitions_batch[:, :, 1]], axis=0)
  
              # This could be done prior to running the batch for a tiny speed
              # boost.
              t_idxs = list(range(FLAGS.num_samples))
              random.shuffle(t_idxs)
              t_idxs = sorted(t_idxs[:FLAGS.num_samples])
              for t_idx in t_idxs:
                  log = log_entry.rl_sampling.add()
                  gold = transitions_batch[t_idx]
                  pred_tr = tr_transitions_per_example[t_idx]
                  pred_ev = ev_transitions_per_example[t_idx]
                  strength_tr = sparks(
                      [1] + tr_strength[t_idx].tolist(), dec_str)
                  strength_ev = sparks(
                      [1] + ev_strength[t_idx].tolist(), dec_str)
                  _, crossing = evalb.crossing(gold, pred)
                  log.t_idx = t_idx
                  log.crossing = crossing
                  log.gold_lb = "".join(map(str, gold))
                  log.pred_tr = "".join(map(str, pred_tr))
                  log.pred_ev = "".join(map(str, pred_ev))
                  log.strg_tr = strength_tr[1:]
                  log.strg_ev = strength_ev[1:]
  
          if trainer.step > 0 and trainer.step % FLAGS.eval_interval_steps == 0:
              should_log = True
              for index, eval_set in enumerate(eval_iterators):
                  acc, _ = evaluate(
                      FLAGS, model, eval_set, log_entry, logger, trainer, eval_index=index, vocabulary=vocabulary, show_sample=True)
                  if  index == 0:
                      trainer.new_dev_accuracy(acc)
  
              progress_bar.reset()
  
          if trainer.step > FLAGS.ckpt_step and trainer.step % FLAGS.ckpt_interval_steps == 0:
              should_log = True
              trainer.checkpoint()
  
          if should_log:
              logger.LogEntry(log_entry)
  
          progress_bar.step(i=(trainer.step % FLAGS.statistics_interval_steps) + 1,
                            total=FLAGS.statistics_interval_steps)
  
  
  def run(only_forward=False):
      logger = afs_safe_logger.ProtoLogger(log_path(FLAGS),
                                           print_formatter=create_log_formatter(
                                               True, True),
                                           write_proto=FLAGS.write_proto_to_log)
      header = pb.SpinnHeader()
  
      data_manager = get_data_manager(FLAGS.data_type)
  
      logger.Log("Flag Values:\n" +
                 json.dumps(FLAGS.FlagValuesDict(), indent=4, sort_keys=True))
  
      # Get Data and Embeddings
      vocabulary, initial_embeddings, training_data_iter, eval_iterators, training_data_length = \
          load_data_and_embeddings(FLAGS, data_manager, logger,
                                   FLAGS.training_data_path, FLAGS.eval_data_path)
  
      # Build model.
      vocab_size = len(vocabulary)
      num_classes = len(set(data_manager.LABEL_MAP.values()))
  
      model = init_model(
          FLAGS,
          logger,
          initial_embeddings,
          vocab_size,
          num_classes,
          data_manager,
          header)
      time_to_wait_to_lower_lr = min(10000, int(training_data_length / FLAGS.batch_size))
      trainer = ModelTrainer(model, logger, time_to_wait_to_lower_lr, vocabulary, FLAGS) 
      
      header.start_step = trainer.step
      header.start_time = int(time.time())
  
      # Do an evaluation-only run.
      logger.LogHeader(header)  # Start log_entry logging.
      if only_forward:
          log_entry = pb.SpinnEntry()
          for index, eval_set in enumerate(eval_iterators):
              log_entry.Clear()
              acc = evaluate(
                  FLAGS,
                  model,
                  eval_set,
                  log_entry,
                  logger,
                  trainer,
                  vocabulary,
                  show_sample=True,
                  eval_index=index)
              print(log_entry)
              logger.LogEntry(log_entry)
      else:
          train_loop(
              FLAGS,
              model,
              trainer,
              training_data_iter,
              eval_iterators,
              logger,
              vocabulary)
  
  
  if __name__ == '__main__':
      get_flags()
  
      # Parse command line flags.
      FLAGS(sys.argv)
  
      flag_defaults(FLAGS)
  
      if FLAGS.model_type != "RLSPINN":
          raise Exception("Reinforce is only implemented for RLSPINN.")
  
      run(only_forward=FLAGS.expanded_eval_only_mode)
  