Add paragraph reconstruction experiment code

Author: ymym3412

datasets.py

@@ -4,6 +4,67 @@
 from tqdm import tqdm
 
 from collections import Counter
+from copy import deepcopy
+
+
+def load_hotel_review_data(path, sentence_len):
+    """
+    Load Hotel Reviews data from pickle distributed in https://drive.google.com/file/d/0B52eYWrYWqIpQzhBNkVxaV9mMjQ/view
+    This file is published in https://github.com/dreasysnail/textCNN_public
+    
+    :param path: pickle path
+    :return: 
+    """
+    import _pickle as cPickle
+    with open(path, "rb") as f:
+        data = cPickle.load(f, encoding="latin1")
+
+    train_data, test_data = HotelReviewsDataset(data[0], deepcopy(data[2]), deepcopy(data[3]), sentence_len, transform=ToTensor()), \
+                             HotelReviewsDataset(data[1], deepcopy(data[2]), deepcopy(data[3]), sentence_len, transform=ToTensor())
+    return train_data, test_data
+
+
+class HotelReviewsDataset(Dataset):
+    """
+    Hotel Reviews Dataset
+    """
+    def __init__(self, data_list, word2index, index2word, sentence_len, transform=None):
+        self.word2index = word2index
+        self.index2word = index2word
+        self.n_words = len(self.word2index)
+        self.data = data_list
+        self.sentence_len = sentence_len
+        self.transform = transform
+        self.word2index["<PAD>"] = self.n_words
+        self.index2word[self.n_words] = "<PAD>"
+        self.n_words += 1
+        print(self.index2word)
+        temp_list = []
+        for sentence in tqdm(self.data):
+            if len(sentence) > self.sentence_len:
+                # truncate sentence if sentence length is longer than `sentence_len`
+                temp_list.append(np.array(sentence[:self.sentence_len]))
+            else:
+                # pad sentence  with '<PAD>' token if sentence length is shorter than `sentence_len`
+                sent_array = np.lib.pad(np.array(sentence),
+                                        (0, self.sentence_len - len(sentence)),
+                                        "constant",
+                                        constant_values=(self.n_words-1, self.n_words-1))
+                temp_list.append(sent_array)
+        self.data = np.array(temp_list, dtype=np.int32)
+
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        data = self.data[idx]
+        if self.transform:
+            data = self.transform(data)
+        return data
+
+    def vocab_lennght(self):
+        return len(self.word2index)
 
 
 class TextClassificationDataset(Dataset):
@@ -81,7 +142,8 @@ def __getitem__(self, idx):
         sample = {"sentence": sentence, "label": label}
 
         if self.transform:
-            sample = self.transform(sample)
+            sample = {"sentence": self.transform(sample["sentence"]),
+                      "label": self.transform(sample["label"])}
 
         return sample
 
@@ -91,8 +153,5 @@ def vocab_length(self):
 
 class ToTensor(object):
     """Convert ndarrays in sample to Tensors."""
-
-    def __call__(self, sample):
-        sentence, label = sample["sentence"], sample['label']
-        return {'sentence': torch.from_numpy(sentence).type(torch.LongTensor),
-                'label': torch.from_numpy(label).type(torch.LongTensor)}
+    def __call__(self, data):
+        return torch.from_numpy(data).type(torch.LongTensor)

main.py renamed to main_classification.py

@@ -4,7 +4,7 @@
 
 import model
 from datasets import TextClassificationDataset, ToTensor
-from train import train
+from train import train_classification
 
 import argparse
 
@@ -17,7 +17,7 @@ def main():
     parser.add_argument('-batch_size', type=int, default=64, help='batch size for training')
     parser.add_argument('-lr_decay_interval', type=int, default=20,
                         help='how many epochs to wait before decrease learning rate')
-    parser.add_argument('-log_interval', type=int, default=256,
+    parser.add_argument('-log_interval', type=int, default=16,
                         help='how many steps to wait before logging training status')
     parser.add_argument('-test_interval', type=int, default=100,
                         help='how many steps to wait before testing')
@@ -68,7 +68,7 @@ def main():
         decoder = torch.load(args.dec_snapshot)
         mlp = torch.load(args.mlp_snapshot)
 
-    train(data_loader, data_loader, encoder, decoder, mlp, args)
+    train_classification(data_loader, data_loader, encoder, decoder, mlp, args)
 
 if __name__ == '__main__':
     main()

main_reconstruction.py

@@ -0,0 +1,62 @@
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+
+import model
+from datasets import TextClassificationDataset, ToTensor, load_hotel_review_data
+from train import train_reconstruction
+
+import argparse
+
+
+def main():
+    parser = argparse.ArgumentParser(description='text convolution-deconvolution auto-encoder model')
+    # learning
+    parser.add_argument('-lr', type=float, default=0.001, help='initial learning rate')
+    parser.add_argument('-epochs', type=int, default=10, help='number of epochs for train')
+    parser.add_argument('-batch_size', type=int, default=32, help='batch size for training')
+    parser.add_argument('-lr_decay_interval', type=int, default=4,
+                        help='how many epochs to wait before decrease learning rate')
+    parser.add_argument('-log_interval', type=int, default=256,
+                        help='how many steps to wait before logging training status')
+    parser.add_argument('-test_interval', type=int, default=10,
+                        help='how many epochs to wait before testing')
+    parser.add_argument('-save_interval', type=int, default=2,
+                        help='how many epochs to wait before saving')
+    parser.add_argument('-save_dir', type=str, default='rec_snapshot', help='where to save the snapshot')
+    # data
+    parser.add_argument('-data_path', type=str, help='data path')
+    parser.add_argument('-shuffle', default=False, help='shuffle data every epoch')
+    parser.add_argument('-sentence_len', type=int, default=210, help='how many tokens in a sentence')
+    # model
+    parser.add_argument('-embed_dim', type=int, default=300, help='number of embedding dimension')
+    parser.add_argument('-kernel_sizes', type=int, default=2,
+                        help='kernel size to use for convolution')
+    parser.add_argument('-tau', type=float, default=0.01, help='temperature parameter')
+    parser.add_argument('-use_cuda', action='store_true', default=True, help='whether using cuda')
+    # option
+    parser.add_argument('-enc_snapshot', type=str, default=None, help='filename of encoder snapshot ')
+    parser.add_argument('-dec_snapshot', type=str, default=None, help='filename of decoder snapshot ')
+    args = parser.parse_args()
+
+    train_data, test_data = load_hotel_review_data(args.data_path, args.sentence_len)
+    train_loader, test_loader = DataLoader(train_data, batch_size=args.batch_size, shuffle=args.shuffle),\
+                                  DataLoader(test_data, batch_size=args.batch_size, shuffle=args.shuffle)
+
+    k = args.embed_dim
+    v = train_data.vocab_lennght()
+    if args.enc_snapshot is None or args.dec_snapshot is None or args.mlp_snapshot is None:
+        print("Start from initial")
+        embedding = nn.Embedding(v, k, max_norm=1.0, norm_type=2.0)
+
+        encoder = model.ConvolutionEncoder(embedding)
+        decoder = model.DeconvolutionDecoder(embedding, args.tau)
+    else:
+        print("Restart from snapshot")
+        encoder = torch.load(args.enc_snapshot)
+        decoder = torch.load(args.dec_snapshot)
+
+    train_reconstruction(train_loader, test_loader, encoder, decoder, args)
+
+if __name__ == '__main__':
+    main()

model.py

@@ -46,7 +46,7 @@ def __init__(self, embedding, tau):
         self.bn1 = nn.BatchNorm2d(600)
         self.deconvs2 = nn.ConvTranspose2d(600, 300, (2, 1), stride=2)
         self.bn2 = nn.BatchNorm2d(300)
-        self.deconvs3 = nn.ConvTranspose2d(300, 1, (2+1, self.embed.weight.size()[1]), stride=2)
+        self.deconvs3 = nn.ConvTranspose2d(300, 1, (2+2, self.embed.weight.size()[1]), stride=2)
 
         # weight initialize for conv_transpose layer
         for m in self.modules():

train.py

@@ -2,10 +2,13 @@
 from torch.autograd import Variable
 import torch.nn.functional as F
 import pickle
+from sumeval.metrics.rouge import RougeCalculator
+from sumeval.metrics.bleu import BLEUCalculator
+from hyperdash import Experiment
 
 import util
 
-def train(data_loader, dev_iter, encoder, decoder, mlp, args):
+def train_classification(data_loader, dev_iter, encoder, decoder, mlp, args):
     lr = args.lr
     encoder_opt = torch.optim.Adam(encoder.parameters(), lr=lr)
     decoder_opt = torch.optim.Adam(decoder.parameters(), lr=lr)
@@ -53,14 +56,13 @@ def train(data_loader, dev_iter, encoder, decoder, mlp, args):
                 input_label = target[0]
                 single_data = prob[0]
                 _, predict_index = torch.max(single_data, 1)
-                input_sentence = util.transform_id2word(input_data, data_loader.dataset.index2word)
-                predict_sentence = util.transform_id2word(predict_index, data_loader.dataset.index2word)
+                input_sentence = util.transform_id2word(input_data.data, data_loader.dataset.index2word, lang="ja")
+                predict_sentence = util.transform_id2word(predict_index.data, data_loader.dataset.index2word, lang="ja")
                 print("Input Sentence:")
                 print(input_sentence)
                 print("Output Sentence:")
                 print(predict_sentence)
-                eval_model(encoder, mlp, input_data, input_label)
-
+                eval_classification(encoder, mlp, input_data, input_label)
 
         if epoch % args.lr_decay_interval == 0:
             # decrease learning rate
@@ -91,13 +93,87 @@ def train(data_loader, dev_iter, encoder, decoder, mlp, args):
     print("Finish!!!")
 
 
+def train_reconstruction(train_loader, test_loader, encoder, decoder, args):
+    lr = args.lr
+    encoder_opt = torch.optim.Adam(encoder.parameters(), lr=lr)
+    decoder_opt = torch.optim.Adam(decoder.parameters(), lr=lr)
+
+    encoder.train()
+    decoder.train()
+    steps = 0
+    for epoch in range(1, args.epochs+1):
+        print("=======Epoch========")
+        print(epoch)
+        for batch in train_loader:
+            feature = Variable(batch)
+            if args.use_cuda:
+                encoder.cuda()
+                decoder.cuda()
+                feature = feature.cuda()
+
+            encoder_opt.zero_grad()
+            decoder_opt.zero_grad()
+
+            h = encoder(feature)
+            prob = decoder(h)
+            reconstruction_loss = compute_cross_entropy(prob, feature)
+            reconstruction_loss.backward()
+            encoder_opt.step()
+            decoder_opt.step()
+
+            steps += 1
+            print("Epoch: {}".format(epoch))
+            print("Steps: {}".format(steps))
+            print("Loss: {}".format(reconstruction_loss.data[0]))
+            # check reconstructed sentence
+            if steps % args.log_interval == 0:
+                print("Test!!")
+                input_data = feature[0]
+                single_data = prob[0]
+                _, predict_index = torch.max(single_data, 1)
+                input_sentence = util.transform_id2word(input_data.data, train_loader.dataset.index2word, lang="en")
+                predict_sentence = util.transform_id2word(predict_index.data, train_loader.dataset.index2word, lang="en")
+                print("Input Sentence:")
+                print(input_sentence)
+                print("Output Sentence:")
+                print(predict_sentence)
+
+        if epoch % args.test_interval == 0:
+            eval_reconstruction(encoder, decoder, test_loader, args)
+
+
+        if epoch % args.lr_decay_interval == 0:
+            # decrease learning rate
+            lr = lr / 5
+            encoder_opt = torch.optim.Adam(encoder.parameters(), lr=lr)
+            decoder_opt = torch.optim.Adam(decoder.parameters(), lr=lr)
+            encoder.train()
+            decoder.train()
+
+        if epoch % args.save_interval == 0:
+            util.save_models(encoder, args.save_dir, "encoder", steps)
+            util.save_models(decoder, args.save_dir, "decoder", steps)
+
+    # finalization
+    # save vocabulary
+    with open("word2index", "wb") as w2i, open("index2word", "wb") as i2w:
+        pickle.dump(train_loader.dataset.word2index, w2i)
+        pickle.dump(train_loader.dataset.index2word, i2w)
+
+    # save models
+    util.save_models(encoder, args.save_dir, "encoder", "final")
+    util.save_models(decoder, args.save_dir, "decoder", "final")
+
+    print("Finish!!!")
+
+
 def compute_cross_entropy(log_prob, target):
     # compute reconstruction loss using cross entropy
     loss = [F.nll_loss(sentence_emb_matrix, word_ids, size_average=False) for sentence_emb_matrix, word_ids in zip(log_prob, target)]
     average_loss = sum([torch.sum(l) for l in loss]) / log_prob.size()[0]
     return average_loss
 
-def eval_model(encoder, mlp, feature, label):
+def eval_classification(encoder, mlp, feature, label):
     encoder.eval()
     mlp.eval()
     h = encoder(feature)
@@ -110,3 +186,45 @@ def eval_model(encoder, mlp, feature, label):
     encoder.train()
     mlp.train()
 
+
+def eval_reconstruction(encoder, decoder, data_iter, args):
+    print("Eval")
+    encoder.eval()
+    decoder.eval()
+    avg_loss = 0
+    rouge_1 = 0.0
+    rouge_2 = 0.0
+    index2word = data_iter.dataset.index2word
+    for batch in data_iter:
+        feature = Variable(batch)
+        if args.use_cuda:
+            feature = feature.cuda()
+        h = encoder(feature)
+        prob = decoder(h)
+        _, predict_index = torch.max(prob, 2)
+        original_sentences = [util.transform_id2word(sentence, index2word, "en") for sentence in batch]
+        predict_sentences = [util.transform_id2word(sentence, index2word, "en") for sentence in predict_index.data]
+        r1, r2 = calc_rouge(original_sentences, predict_sentences)
+        rouge_1 += r1
+        rouge_2 += r2
+        reconstruction_loss = compute_cross_entropy(prob, feature)
+        avg_loss += reconstruction_loss.data[0]
+    avg_loss = avg_loss / len(data_iter.dataset)
+    rouge_1 = rouge_1 / len(data_iter.dataset)
+    rouge_2 = rouge_2 / len(data_iter.dataset)
+    print("Evaluation - loss: {}  Rouge1: {}    Rouge2: {}".format(avg_loss, rouge_1, rouge_2))
+    encoder.train()
+    decoder.train()
+
+def calc_rouge(original_sentences, predict_sentences):
+    rouge_1 = 0.0
+    rouge_2 = 0.0
+    for original, predict in zip(original_sentences, predict_sentences):
+        # Remove padding
+        original, predict = original.replace("<PAD>", "").strip(), predict.replace("<PAD>", "").strip()
+        rouge = RougeCalculator(stopwords=True, lang="en")
+        r1 = rouge.rouge_1(summary=predict, references=original)
+        r2 = rouge.rouge_2(summary=predict, references=original)
+        rouge_1 += r1
+        rouge_2 += r2
+    return rouge_1, rouge_2

util.py

@@ -2,8 +2,11 @@
 import math
 import os
 
-def transform_id2word(index, id2word):
-    return "".join([id2word[idx.data[0]] for idx in index])
+def transform_id2word(index, id2word, lang):
+    if lang == "ja":
+        return "".join([id2word[idx] for idx in index])
+    else:
+        return " ".join([id2word[idx] for idx in index])
 
 def sigmoid_annealing_schedule(step, max_step, param_init=1.0, param_final=0.01, gain=0.3):
     return ((param_init - param_final) / (1 + math.exp(gain * (step - (max_step / 2))))) + param_final