error_analysis.py

# Copyright 2016 Stanford University
#
# 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.
# ==============================================================================

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import math
import os
import random
import sys
import time
import random
import string

import numpy as np
from six.moves import xrange
import tensorflow as tf
import csv
import itertools

import kenlm

import nlc_model
import nlc_data
from util import pair_iter

tf.app.flags.DEFINE_float("learning_rate", 0.001, "Learning rate.")
tf.app.flags.DEFINE_float("learning_rate_decay_factor", 0.95, "Learning rate decays by this much.")
tf.app.flags.DEFINE_float("max_gradient_norm", 5.0, "Clip gradients to this norm.")
tf.app.flags.DEFINE_float("dropout", 0.1, "Fraction of units randomly dropped on non-recurrent connections.")
tf.app.flags.DEFINE_integer("batch_size", 128, "Batch size to use during training.")
tf.app.flags.DEFINE_integer("epochs", 0, "Number of epochs to train.")
tf.app.flags.DEFINE_integer("size", 400, "Size of each model layer.")
tf.app.flags.DEFINE_integer("num_layers", 3, "Number of layers in the model.")
tf.app.flags.DEFINE_integer("max_vocab_size", 40000, "Vocabulary size limit.")
tf.app.flags.DEFINE_integer("max_seq_len", 200, "Maximum sequence length.")
tf.app.flags.DEFINE_string("data_dir", "/tmp", "Data directory")
tf.app.flags.DEFINE_string("train_dir", "/tmp", "Training directory.")
tf.app.flags.DEFINE_string("tokenizer", "CHAR", "Set to WORD to train word level model.")
tf.app.flags.DEFINE_integer("beam_size", 8, "Size of beam.")
tf.app.flags.DEFINE_string("lmfile", None, "arpa file of the language model.")
tf.app.flags.DEFINE_float("alpha", 0.3, "Language model relative weight.")
tf.app.flags.DEFINE_boolean("dev", False, "generate dev outputs with batch_decode")
tf.app.flags.DEFINE_boolean("sweep", False, "sweep all alpha rates with dev turned on")
tf.app.flags.DEFINE_boolean("score", False, "generate csv with language model scores on target and generated.")

FLAGS = tf.app.flags.FLAGS
reverse_vocab, vocab = None, None
lm = None

def create_model(session, vocab_size, forward_only):
  model = nlc_model.NLCModel(
      vocab_size, FLAGS.size, FLAGS.num_layers, FLAGS.max_gradient_norm, FLAGS.batch_size,
      FLAGS.learning_rate, FLAGS.learning_rate_decay_factor, FLAGS.dropout,
      forward_only=forward_only)
  ckpt = tf.train.get_checkpoint_state(FLAGS.train_dir)
  if ckpt and tf.gfile.Exists(ckpt.model_checkpoint_path):
    print("Reading model parameters from %s" % ckpt.model_checkpoint_path)
    model.saver.restore(session, ckpt.model_checkpoint_path)
  else:
    print("Created model with fresh parameters.")
    session.run(tf.initialize_all_variables())
  return model


def get_tokenizer(FLAGS):
  tokenizer = nlc_data.char_tokenizer if FLAGS.tokenizer.lower() == 'char' else nlc_data.basic_tokenizer
  return tokenizer

def detokenize(sents, reverse_vocab):
  # TODO: char vs word
  def detok_sent(sent):
    outsent = ''
    for t in sent:
      if t >= len(nlc_data._START_VOCAB):
        outsent += reverse_vocab[t]
    return outsent
  return [detok_sent(s) for s in sents]

def network_score(model, sess, encoder_output, target_tokens):
  score = 0.0
  states = None
  for (feed, pick) in zip(list(target_tokens)[:-1], list(target_tokens)[1:]):
    scores, _, states = model.decode(sess, encoder_output, np.array([feed]), None, states)
    score += float(scores[0, 0, pick])
  return score

def detokenize_tgt(toks, reverse_vocab):
  outsent = ''
  for i in range(toks.shape[0]):
    if toks[i] >= len(nlc_data._START_VOCAB) and toks[i] != nlc_data._PAD:
      outsent += reverse_vocab[toks[i][0]]
  return outsent

def lm_rank(strs, probs):
  if lm is None:
    return strs[0]
  a = FLAGS.alpha
  rescores = [(1-a)*p + a*lm.score(s) for (s, p) in zip(strs, probs)]
  rerank = [rs[0] for rs in sorted(enumerate(rescores), key=lambda x:x[1])]
  return strs[rerank[-1]]

def lm_rank_score(strs, probs):
  """

  :param strs: candidates generated by beam search
  :param target: target sentence
  :param probs:
  :return:
  """
  if lm is None:
    return strs[0]
  a = FLAGS.alpha
  lmscores = [lm.score(s) for s in strs]
  rescores = [(1 - a) * p + a * l for (l, p) in zip(lmscores, probs)]
  rerank = [rs[0] for rs in sorted(enumerate(rescores), key=lambda x: x[1])]
  generated = strs[rerank[-1]]
  lm_score = lmscores[rerank[-1]]
  nw_score = probs[rerank[-1]]
  score = rescores[rerank[-1]]
  return generated, score, nw_score, lm_score

def decode_beam(model, sess, encoder_output, max_beam_size):
  toks, probs = model.decode_beam(sess, encoder_output, beam_size=max_beam_size)
  return toks.tolist(), probs.tolist()

def setup_batch_decode(sess):
  # decode for dev-sets, in batches
  global reverse_vocab, vocab, lm

  if FLAGS.lmfile is not None:
    print("Loading Language model from %s" % FLAGS.lmfile)
    lm = kenlm.LanguageModel(FLAGS.lmfile)

  print("Preparing NLC data in %s" % FLAGS.data_dir)

  x_train, y_train, x_dev, y_dev, vocab_path = nlc_data.prepare_nlc_data(
    FLAGS.data_dir + '/' + FLAGS.tokenizer.lower(), FLAGS.max_vocab_size,
    tokenizer=get_tokenizer(FLAGS))
  vocab, reverse_vocab = nlc_data.initialize_vocabulary(vocab_path)
  vocab_size = len(vocab)
  print("Vocabulary size: %d" % vocab_size)

  print("Creating %d layers of %d units." % (FLAGS.num_layers, FLAGS.size))
  model = create_model(sess, vocab_size, False)

  return model, x_dev, y_dev


def batch_decode(model, sess, x_dev, y_dev, alpha):

    error_source = [];
    error_target = [];
    error_generated = [];
    generated_score = [];
    generated_lm_score = [];
    generated_nw_score = [];
    target_lm_score = [];
    target_nw_score = [];

    for source_tokens, source_mask, target_tokens, target_mask in pair_iter(x_dev, y_dev, 1,
                                                                            FLAGS.num_layers):
      src_sent = detokenize_tgt(source_tokens, reverse_vocab)
      tgt_sent = detokenize_tgt(target_tokens, reverse_vocab)

      # Encode
      encoder_output = model.encode(sess, source_tokens, source_mask)
      # Decode
      beam_toks, probs = decode_beam(model, sess, encoder_output, FLAGS.beam_size)
      # De-tokenize
      beam_strs = detokenize(beam_toks, reverse_vocab)
      tgt_nw_score = network_score(model, sess, encoder_output, target_tokens)
      print("Network score: %f" % tgt_nw_score)
      # Language Model ranking
      if not FLAGS.score:
        best_str = lm_rank(beam_strs, probs)
      else:
        best_str, rerank_score, nw_score, lm_score = lm_rank_score(beam_strs, probs)
        tgt_lm_score = lm.score(tgt_sent)

      # see if this is too stupid, or doesn't work at all
      error_source.append(src_sent)
      error_target.append(tgt_sent)
      error_generated.append(best_str)
      if FLAGS.score:
        target_lm_score.append(tgt_lm_score)
        target_nw_score.append(tgt_nw_score)
        generated_score.append(rerank_score)
        generated_nw_score.append(nw_score)
        generated_lm_score.append(lm_score)

    print("outputting in csv file...")

    # dump it out in train_dir
    with open(FLAGS.train_dir + "/err_analysis/" + "err_val_alpha_" + str(alpha) + ".csv", 'wb') as f:
      wrt = csv.writer(f)
      wrt.writerow(['Bad Input', 'Ground Truth', 'Network Score', 'LM Score', 'Generated Hypothesis', 'Combined Score', 'Network Score', 'LM Score'])
      if not FLAGS.score:
        for s, t, g in itertools.izip(error_source, error_target, error_generated):
          wrt.writerow([s, t, g])  # source, correct target, wrong target
      else:
        for s, t, tns, tls, g, gs, gns, gls in itertools.izip(error_source, error_target, target_nw_score, target_lm_score, error_generated, generated_score, generated_nw_score, generated_lm_score):
          wrt.writerow([s, t, tns, tls, g, gs, gns, gls]) 

    print("err_val_alpha_" + str(alpha) + ".csv" + "file finished")


def main(_):
  with tf.Session() as sess:
    if FLAGS.sweep:
      alpha = [0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1]
      model, x_dev, y_dev = setup_batch_decode(sess)
      for a in alpha:
        FLAGS.alpha = a
        print("ranking with alpha = " + str(FLAGS.alpha))
        batch_decode(model, sess, x_dev, y_dev, a)
    else:
      model, x_dev, y_dev = setup_batch_decode(sess)
      print("ranking with alpha = " + str(FLAGS.alpha))
      batch_decode(model, sess, x_dev, y_dev, FLAGS.alpha)

if __name__ == "__main__":
  tf.app.run()