train_styleptb.py

#!/usr/bin496/env python3
import sys
import os

import argparse
import json
import random
import shutil
import copy
import pickle
import torch
from torch import cuda
import numpy as np
import time
import logging
from tokenizer import Tokenizer

from utils import *
from torch.nn.init import xavier_uniform_
from torch.nn.utils.rnn import pad_sequence


parser = argparse.ArgumentParser()

parser.add_argument('--train_file', default='data/StylePTB/ATP/train.tsv')
parser.add_argument('--dev_file', default='data/StylePTB/ATP/valid.tsv')
parser.add_argument('--save_path', default='model.pt', help='where to save the model')
parser.add_argument('--min_freq', default=3, type=int)
parser.add_argument('--sent_max_length_x', default=20, type=int)
parser.add_argument('--sent_max_length_y', default=20, type=int)
# Encoder 
parser.add_argument('--enc_dim', default=512, type=int)
parser.add_argument('--enc_layers', default=1, type=int)
parser.add_argument('--enc_dropout', default=0.0, type=float)
# Decoder 
parser.add_argument('--pt_states', default=0, type=int)
parser.add_argument('--nt_states', default=0, type=int)
parser.add_argument('--src_pt_states', default=8, type=int)
parser.add_argument('--src_nt_states', default=8, type=int)
parser.add_argument('--dec_dim', default=512, type=int)
parser.add_argument('--dec_char_dim', default=512, type=int)
parser.add_argument('--dec_dropout', default=0.0, type=float)
parser.add_argument('--dec_layers', default=3, type=int)
parser.add_argument('--dec_nt_span_min', default=2, type=int)
parser.add_argument('--dec_nt_span_max', default=100, type=int)
parser.add_argument('--dec_pt_span_min', default=1, type=int)
parser.add_argument('--dec_pt_span_max', default=1, type=int)
parser.add_argument('--rule_constraint_type', default=1, type=int)
parser.add_argument('--use_copy', default=1, type=int)
# Parser 
parser.add_argument('--parser_pt_states', default=20, type=int)
parser.add_argument('--parser_nt_states', default=20, type=int)
parser.add_argument('--parser_dim', default=256, type=int)
# Optimization
parser.add_argument('--num_epochs', default=15, type=int, help='number of training epochs')
parser.add_argument('--lr', default=5e-4, type=float, help='starting learning rate')
parser.add_argument('--weight_decay', default=1e-5, type=float, help='l2 weight decay')
parser.add_argument('--max_grad_norm', default=3, type=float, help='gradient clipping parameter')
parser.add_argument('--beta1', default=0.75, type=float, help='beta1 for adam')
parser.add_argument('--beta2', default=0.999, type=float, help='beta2 for adam')
parser.add_argument('--gpu', default=0, type=int, help='which gpu to use')
parser.add_argument('--seed', default=3435, type=int, help='random seed')
parser.add_argument('--print_every', type=int, default=500, help='print stats after N examples')
parser.add_argument('--print_trees', type=int, default=1, help='print trees')
parser.add_argument('--eval_every', type=int, default=500, help='eval on dev set after N examples')
parser.add_argument('--update_every', type=int, default=4, help='grad update after N examples')

def get_data(data_file):
  data = []
  for d in open(data_file):
    src, tgt = d.split("\t")
    if ";" in src:
      src, emph = src.strip().split(";")
      emph = emph.strip()
      src = src.strip().split()
      emph_mask = []
      for w in src:
        if w == emph:
          emph_mask.append(1)
        else:
          emph_mask.append(0)
      data.append({"src": src, "tgt": tgt.strip().split(), "emph_mask": emph_mask})
    else:
      data.append({"src": src.strip().split(), "tgt": tgt.strip().split()})
  return data

def main(args):
  random.seed(args.seed)
  np.random.seed(args.seed)
  torch.manual_seed(args.seed)
  torch.cuda.manual_seed(args.seed)
  cuda.set_device(args.gpu)
  device = torch.device("cuda:"+str(args.gpu))
  train_data = get_data(args.train_file)
  val_data = get_data(args.dev_file)

  tokenizer = Tokenizer()
  tokenizer.train([d["src"] for d in train_data if len(d["src"]) <= args.sent_max_length_x] + \
                  [d["tgt"] for d in train_data if len(d["tgt"]) <= args.sent_max_length_y],
                  False, args.min_freq)
  from models import BinaryTreeLSTM as Encoder
  from models import NeuralQCFG as Decoder
  from models import NeuralPCFG as Parser

  encoder = Encoder(vocab = len(tokenizer.vocab2idx),
                    dim = args.enc_dim,
                    dropout = args.enc_dropout,
                    layers = args.enc_layers,
                    token_type_emb = torch.nn.Embedding(2, args.enc_dim))
  decoder = Decoder(vocab = len(tokenizer.vocab2idx),
                    dim = args.dec_dim,
                    num_layers = args.dec_layers,
                    pt_states = args.pt_states,
                    nt_states = args.nt_states,
                    src_dim = args.enc_dim,
                    src_pt_states = args.src_pt_states,
                    src_nt_states = args.src_nt_states, 
                    dropout = args.dec_dropout,
                    rule_constraint_type = args.rule_constraint_type,
                    use_copy = args.use_copy == 1,
                    tokenizer = tokenizer, 
                    nt_span_range = [args.dec_nt_span_min, args.dec_nt_span_max],
                    pt_span_range = [args.dec_pt_span_min, args.dec_pt_span_max])
  parser = Parser(vocab = len(tokenizer.vocab2idx),
                      dim = args.parser_dim,
                      nt_states = args.parser_nt_states,
                      pt_states = args.parser_pt_states)
  model = torch.nn.ModuleList([encoder, decoder, parser])
  for name, param in model.named_parameters():    
    if param.dim() > 1:
      xavier_uniform_(param)
  optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, 
                               betas = (args.beta1, args.beta2), 
                               weight_decay = args.weight_decay)
  best_val_ppl = 1e5
  epoch = 0
  b = 0
  model.to(device)
  model.train()
  while epoch < args.num_epochs:
    start_time = time.time()
    epoch += 1  
    print('Starting epoch: %d' % epoch)
    train_nll = 0.
    src_nll = 0.
    num_sents = 0.
    num_src_words = 0.
    num_words = 0.
    random.shuffle(train_data)
    for d in train_data:
      if len(d["src"]) > args.sent_max_length_x or len(d["tgt"]) > args.sent_max_length_y or \
         len(d["src"]) < 2 or len(d["tgt"]) < 2:
          continue
      b += 1
      x = [d["src"]]
      y = [d["tgt"]]
      x_tensor, _, _ = tokenizer.convert_batch(x)
      y_tensor, _, _ = tokenizer.convert_batch(y)
      x_tensor, y_tensor = x_tensor.to(device), y_tensor.to(device)
      emph_mask = torch.LongTensor(d["emph_mask"]).to(device) if "emph_mask" in d else None      
      x_lengths = torch.Tensor([len(d["src"])]).long().to(device)
      y_lengths = torch.Tensor([len(d["tgt"])]).long().to(device)
      parse_sample, parse_argmax, parse_log_prob, parse_actions, parse_nll = parser(
        x_tensor, x_lengths)
      node_features, node_spans = encoder(x_tensor, x_lengths, spans = parse_sample,
                                          token_type = emph_mask)
      new_spans = []
      for span, x_str in zip(node_spans, x):
        new_span = []
        for s in span:
          new_span.append([s[0], s[1], x_str[s[0]:s[1]+1]])
        new_spans.append(new_span)
      node_spans = new_spans
      nll = decoder(y_tensor, y_lengths, node_features, node_spans,
                    x_str = y, argmax=False)
      dec_loss = nll.mean()
      (dec_loss / args.update_every).backward()
      train_nll += nll.sum().item()        
      with torch.no_grad():
        node_features_argmax, node_spans_argmax = encoder(x_tensor, x_lengths,
                                                          spans = parse_argmax)
        new_spans = []
        for span, x_str in zip(node_spans_argmax, x):
          new_span = []
          for s in span:
            new_span.append([s[0], s[1], x_str[s[0]:s[1]+1]])
          new_spans.append(new_span)
        node_spans_argmax = new_spans
        nll_argmax = decoder(y_tensor, y_lengths, node_features_argmax, node_spans_argmax, 
                             x_str = y, argmax=False)
        neg_reward = (nll - nll_argmax).detach().item()
      obj = (neg_reward*parse_log_prob).mean() + parse_nll.mean()
      (obj / args.update_every).backward()
      src_nll += parse_nll.sum().item()
      if b % args.update_every == 0:
          torch.nn.utils.clip_grad_norm_(parser.parameters(), args.max_grad_norm)      
          torch.nn.utils.clip_grad_norm_(encoder.parameters(), args.max_grad_norm)      
          torch.nn.utils.clip_grad_norm_(decoder.parameters(), args.max_grad_norm)      
          optimizer.step()
          optimizer.zero_grad()
      num_sents += len(y)
      num_words += y_lengths.sum().item()
      num_src_words += x_lengths.sum().item()
      if b % args.print_every == 0:
        enc_param_norm = sum([p.norm()**2 for p in encoder.parameters()]).item()**0.5
        dec_param_norm = sum([p.norm()**2 for p in decoder.parameters()]).item()**0.5
        parser_param_norm = sum([p.norm()**2 for p in parser.parameters()]).item()**0.5

        log_str = 'Epoch: %d, Batch: %d/%d, |EncParam|: %.4f, |DecParam|: %.4f, ' + \
                  '|ParserParam|: %.4f, LR: %.4f, SrcPPL: %.4f, ' + \
                  'PPL: %.4f, ValPPL: %.4f, ' + \
                  'Throughput: %.2f examples/sec'
        print("-"*80)
        print(log_str %
              (epoch, b, len(train_data),
               enc_param_norm, dec_param_norm, parser_param_norm, 
               args.lr, np.exp(src_nll / num_src_words), 
               np.exp(train_nll / num_words), best_val_ppl, 
               num_sents / (time.time() - start_time)))
        print("-"*80)
        if args.print_trees == 1:
          print("")
          with torch.no_grad():
            y_tree, all_spans, all_spans_node = decoder(
              y_tensor, y_lengths, node_features, node_spans,
              x_str = y, argmax=True)
          x_idx2vocab = tokenizer.idx2vocab
          y_idx2vocab = tokenizer.idx2vocab
          x_str = [x_idx2vocab[idx] for idx in x_tensor[0].tolist()]
          y_str = [y_idx2vocab[idx] for idx in y_tensor[0].tolist()]
          x_length = x_lengths[0].item()
          y_length = y_lengths[0].item()
          print("Source: %s\nTarget: %s" % (" ".join(x_str), " ".join(y_str)))
          print("Source Tree: %s" % get_tree(parse_actions[0], x_str))
          print("")
          action = get_actions(y_tree[0])              
          print("QCFG Tree: %s" % get_tree(action, y_str))
          print("")
          for span, span_node in zip(all_spans[0], all_spans_node[0]):
            if span_node[0] == -1:
              if span[0] == span[1]:
                x_span = "T" + str(span_node[2])
              else:
                x_span = "NT" + str(span_node[2])
            else:
              x_span = " ".join(x_str[span_node[0]:span_node[1]+1])
            y_span = " ".join(y_str[span[0]:span[1]+1])
            if span[0] == span[1]:
              denom = len(decoder.pt_spans[0])
            else:
              denom = len(decoder.nt_spans[0])
            print((y_span, x_span, "N" + str(span[2] // denom)))
        if b % args.eval_every == 0:        
          print('--------------------------------')
          print('Checking validation perf...')    
          val_ppl = eval(val_data, encoder, decoder, parser, device, tokenizer)
          print('--------------------------------')
          if val_ppl < best_val_ppl:
            best_val_ppl = val_ppl
            checkpoint = {
              'args': args.__dict__,
              'encoder': encoder.cpu(),
              'decoder': decoder.cpu(),
              'parser': parser.cpu(),
              'tokenizer': tokenizer,
            }
            print('Saving checkpoint to %s' % args.save_path)
            torch.save(checkpoint, args.save_path)
            model.to(device)

def eval(data, encoder, decoder, parser, device, tokenizer):
  encoder.eval()
  decoder.eval()
  parser.eval()
  num_sents = 0
  num_words = 0
  total_nll = 0.
  b = 0
  for d in data:
    if len(d["src"]) > args.sent_max_length_x or len(d["tgt"]) > args.sent_max_length_y or \
       len(d["src"]) < 2 or len(d["tgt"]) < 2:
        continue
    b += 1
    x = [d["src"]]
    y = [d["tgt"]]
    x_tensor, _, _ = tokenizer.convert_batch(x)
    y_tensor, _, _ = tokenizer.convert_batch(y)
    x_tensor, y_tensor = x_tensor.to(device), y_tensor.to(device)
    emph_mask = torch.LongTensor(d["emph_mask"]).to(device) if "emph_mask" in d else None      
    x_lengths = torch.Tensor([len(d["src"])]).long().to(device)
    y_lengths = torch.Tensor([len(d["tgt"])]).long().to(device)
    parse_sample, parse_argmax, parse_log_prob, parse_actions, parse_nll = parser(
      x_tensor, x_lengths)
    with torch.no_grad():
      node_features, node_spans = encoder(x_tensor, x_lengths, spans = parse_argmax,
                                          token_type = emph_mask)
      new_spans = []
      for span, x_str in zip(node_spans, x):
        new_span = []
        for s in span:
          new_span.append([s[0], s[1], x_str[s[0]:s[1]+1]])
        new_spans.append(new_span)
      node_spans = new_spans
      nll = decoder(y_tensor, y_lengths, node_features, node_spans,
                    x_str = y, argmax=False)
    total_nll += nll.sum().item()
    num_words += y_lengths.sum().item()
  ppl = np.exp(total_nll / num_words)
  print('PPL: %.4f' % ppl)
  encoder.train()
  decoder.train()
  parser.train()
  return ppl

if __name__ == '__main__':
  args = parser.parse_args()
  main(args)