dual learning

Author: JCly-rikiu

data.py

@@ -0,0 +1,63 @@
+import os
+import torch
+import pickle
+
+
+class Dictionary(object):
+    def __init__(self):
+        self.word2idx = {'<unk>': 0}
+        self.idx2word = ['<unk>']
+        self.wordcnt = {}
+
+    def add_word(self, word):
+        if word not in self.word2idx:
+            self.idx2word.append(word)
+            self.word2idx[word] = len(self.idx2word) - 1
+            self.wordcnt[word] = 1
+        else:
+            self.wordcnt[word] = self.wordcnt[word] + 1
+        return self.word2idx[word]
+
+    def getid(self, word, thresh=10):
+        if (word not in self.word2idx) or (self.wordcnt[word] < thresh):
+            return self.word2idx['<unk>']
+        return self.word2idx[word]
+
+    def __len__(self):
+        return len(self.idx2word)
+
+
+class Corpus(object):
+    def __init__(self, path):
+        self.dictionary = Dictionary()
+        self.train = self.tokenize(os.path.join(path, 'train.txt'))
+        self.valid = self.tokenize(os.path.join(path, 'valid.txt'))
+        self.test = self.tokenize(os.path.join(path, 'test.txt'))
+
+        with open(os.path.join(path, 'dict.pkl'), 'wb') as f:
+            pickle.dump(self.dictionary, f)
+
+    def tokenize(self, path):
+        """Tokenizes a text file."""
+        assert os.path.exists(path)
+        # Add words to the dictionary
+        with open(path, 'r') as f:
+            tokens = 0
+            for line in f:
+                words = ['<sos>'] + line.split() + ['<eos>']
+                tokens += len(words)
+                for word in words:
+                    self.dictionary.add_word(word)
+
+        # Tokenize file content
+        with open(path, 'r') as f:
+            ids = torch.LongTensor(tokens)
+            token = 0
+            for line in f:
+                words = ['<sos>'] + line.split() + ['<eos>']
+                for word in words:
+                    ids[token] = self.dictionary.getid(word)
+                    token += 1
+
+        return ids
+

dual.py

@@ -0,0 +1,143 @@
+# -*- coding: utf-8 -*-
+
+import sys
+import torch
+import argparse
+import random
+
+from torch.autograd import Variable
+
+from nmt import read_corpus, data_iter
+from nmt import NMT, to_input_variable
+
+from lm import LMProb
+from lm import model
+
+def dual(args):
+    vocabs = {}
+    opts = {}
+    state_dicts = {}
+    train_srcs = {}
+    lms = {}
+
+    # load model params & training data
+    print('load modelA from [{:s}]'.format(args.modelA_bin), file=sys.stderr)
+    params = torch.load(args.modelA_bin, map_location=lambda storage, loc: storage)
+    vocabs['A'] = params['vocab']
+    opts['A'] = params['args']
+    state_dicts['A'] = params['state_dict']
+    print('load train_srcA from [{:s}]'.format(args.train_srcA), file=sys.stderr)
+    train_srcs['A'] = read_corpus(args.train_srcA, source='src')
+    print('load lmA from [{:s}]'.format(args.lmA), file=sys.stderr)
+    lms['A'] = LMProb(args.lmA, args.lmAdict)
+
+    print('load modelB from [{:s}]'.format(args.modelB_bin), file=sys.stderr)
+    params = torch.load(args.modelB_bin, map_location=lambda storage, loc: storage)
+    vocabs['B'] = params['vocab']
+    opts['B'] = params['args']
+    state_dicts['B'] = params['state_dict']    
+    print('load train_srcB from [{:s}]'.format(args.train_srcB), file=sys.stderr)
+    train_srcs['B'] = read_corpus(args.train_srcB, source='src')
+    print('load lmB from [{:s}]'.format(args.lmB), file=sys.stderr)
+    lms['B'] = LMProb(args.lmB, args.lmBdict)
+
+    models = {}
+    optimizers = {}
+
+    for m in ['A', 'B']:
+        # build model
+        models[m] = NMT(opts[m], vocabs[m])
+        models[m].load_state_dict(state_dicts[m])
+        models[m].train()
+        models[m] = models[m].cuda()
+
+        random.shuffle(train_srcs[m])
+
+        # optimizer
+        optimizers[m] = torch.optim.Adam(models[m].parameters())
+
+    # loss function
+    loss_nll = torch.nn.NLLLoss()
+    loss_ce = torch.nn.CrossEntropyLoss()
+
+    epoch = 0
+    while True:
+        if epoch == 2:
+            break
+        epoch += 1
+        print('start of epoch {:d}'.format(epoch))
+
+        data = {}
+        data['A'] = iter(train_srcs['A'])
+        data['B'] = iter(train_srcs['B'])
+
+        for t in range(0, len(train_srcs['A'])):
+            print('sent', t)
+            for m in ['A', 'B']:
+                lm_probs = []
+
+                NLL_losses = []
+                CE_losses = []
+
+                modelA = models[m]
+                modelB = models[change(m)]
+                lmB = lms[change(m)]
+                optimizerA = optimizers[m]
+                optimizerB = optimizers[change(m)]
+                vocabB = vocabs[change(m)]
+                s = next(data[m])
+
+                hyps = modelA.beam(s, beam_size=5)
+
+                for ids, smid, dist in hyps:
+                    var_ids = torch.autograd.Variable(torch.LongTensor(ids[1:]), requires_grad=False)
+                    NLL_losses.append(loss_nll(dist, var_ids).cpu())
+
+                    lm_probs.append(lmB.get_prob(smid))
+
+                    src_sent_var = to_input_variable([smid], vocabB.src, cuda=True)
+                    tgt_sent_var = to_input_variable([['<s>'] + s + ['</s>']], vocabB.tgt, cuda=True)
+                    src_sent_len = [len(smid)]
+
+                    score = modelB(src_sent_var, src_sent_len, tgt_sent_var[:-1]).squeeze(1)
+
+                    CE_losses.append(loss_ce(score, tgt_sent_var[1:].view(-1)).cpu())
+
+                r1_mean = sum(lm_probs) / len(lm_probs)
+                r1 = [Variable(torch.FloatTensor([p - r1_mean]), requires_grad=False) for p in lm_probs]
+
+                r2_mean = sum(CE_losses) / len(CE_losses)
+                r2 = [Variable(-(l.data - r2_mean.data), requires_grad=False) for l in CE_losses]
+
+                rk = [a + b for a, b in zip(r1, r2)]
+
+                optimizerA.zero_grad()
+                optimizerB.zero_grad()
+
+                torch.mean(torch.cat(NLL_losses) * torch.cat(rk)).backward()
+                torch.mean(torch.cat(CE_losses)).backward()
+
+                optimizerA.step()
+                optimizerB.step()
+
+
+def change(m):
+    if m == 'A':
+        return 'B'
+    else:
+        return 'A'
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('modelA_bin')
+    parser.add_argument('modelB_bin')
+    parser.add_argument('lmA')
+    parser.add_argument('lmAdict')
+    parser.add_argument('lmB')
+    parser.add_argument('lmBdict')
+    parser.add_argument('train_srcA')
+    parser.add_argument('train_srcB')
+    args = parser.parse_args()
+
+    dual(args)
+

model.py

@@ -0,0 +1,47 @@
+import torch.nn as nn
+from torch.autograd import Variable
+
+
+class RNNModel(nn.Module):
+    """Container module with an encoder, a recurrent module, and a decoder."""
+
+    def __init__(self, ntoken, ninp, nhid, nlayers, dropout=0.5, tie_weights=False):
+        super(RNNModel, self).__init__()
+        self.drop = nn.Dropout(dropout)
+        self.encoder = nn.Embedding(ntoken, ninp)
+        self.rnn = nn.GRU(ninp, nhid, nlayers, dropout=dropout)
+        self.decoder = nn.Linear(nhid, ntoken)
+
+        # Optionally tie weights as in:
+        # "Using the Output Embedding to Improve Language Models" (Press & Wolf 2016)
+        # https://arxiv.org/abs/1608.05859
+        # and
+        # "Tying Word Vectors and Word Classifiers: A Loss Framework for Language Modeling" (Inan et al. 2016)
+        # https://arxiv.org/abs/1611.01462
+        if tie_weights:
+            if nhid != ninp:
+                raise ValueError('When using the tied flag, nhid must be equal to emsize')
+            self.decoder.weight = self.encoder.weight
+
+        self.init_weights()
+
+        self.nhid = nhid
+        self.nlayers = nlayers
+
+    def init_weights(self):
+        initrange = 0.1
+        self.encoder.weight.data.uniform_(-initrange, initrange)
+        self.decoder.bias.data.fill_(0)
+        self.decoder.weight.data.uniform_(-initrange, initrange)
+
+    def forward(self, input, hidden):
+        emb = self.drop(self.encoder(input))
+        output, hidden = self.rnn(emb, hidden)
+        output = self.drop(output)
+        decoded = self.decoder(output.view(output.size(0)*output.size(1), output.size(2)))
+        return decoded.view(output.size(0), output.size(1), decoded.size(1)), hidden
+
+    def init_hidden(self, bsz):
+        weight = next(self.parameters()).data
+        return Variable(weight.new(self.nlayers, bsz, self.nhid).zero_())
+

util.py

@@ -0,0 +1,56 @@
+from collections import defaultdict
+import numpy as np
+
+def read_corpus(file_path, source):
+    data = []
+    for line in open(file_path):
+        sent = line.strip().split(' ')
+        # only append <s> and </s> to the target sentence
+        if source == 'tgt':
+            sent = ['<s>'] + sent + ['</s>']
+        data.append(sent)
+
+    return data
+
+
+def batch_slice(data, batch_size, sort=True):
+    batched_data = []
+    batch_num = int(np.ceil(len(data) / float(batch_size)))
+    for i in range(batch_num):
+        cur_batch_size = batch_size if i < batch_num - 1 else len(data) - batch_size * i
+        src_sents = [data[i * batch_size + b][0] for b in range(cur_batch_size)]
+        tgt_sents = [data[i * batch_size + b][1] for b in range(cur_batch_size)]
+
+        if sort:
+            src_ids = sorted(range(cur_batch_size), key=lambda src_id: len(src_sents[src_id]), reverse=True)
+            src_sents = [src_sents[src_id] for src_id in src_ids]
+            tgt_sents = [tgt_sents[src_id] for src_id in src_ids]
+
+        batched_data.append((src_sents, tgt_sents))
+
+    return batched_data
+
+
+def data_iter(data, batch_size, shuffle=True):
+    """
+    randomly permute data, then sort by source length, and partition into batches
+    ensure that the length of source sentences in each batch is decreasing
+    """
+
+    buckets = defaultdict(list)
+    for pair in data:
+        src_sent = pair[0]
+        buckets[len(src_sent)].append(pair)
+
+    batched_data = []
+    for src_len in buckets:
+        tuples = buckets[src_len]
+        if shuffle: np.random.shuffle(tuples)
+        batched_data.extend(batch_slice(tuples, batch_size))
+
+    if shuffle:
+        np.random.shuffle(batched_data)
+
+    for batch in batched_data:
+        yield batch
+