Relation prediction model from Simple QA paper (#20)

Initial implementation of RNNs for relation prediction described by Ture and Jojic: https://arxiv.org/abs/1606.05029
2026-09-09 11:13:20 +08:00 · 2017-04-20 13:29:32 -04:00
parent a67e2d12c4
commit d12a9cb475
12 changed files with 506 additions and 0 deletions
@@ -0,0 +1,2 @@
+.DS_Store
+.idea/
@@ -0,0 +1,4 @@
+datasets/
+resources/
+saved_checkpoints/
+__pycache__/
@@ -0,0 +1,34 @@
+Setup:
+1. Create 3 directories under "simple_qa_rnn" - "resources", "datasets", "saved_checkpoints" 
+2. Download the SimpleQA dataset from [here](https://github.com/castorini/data) and put it under the "datasets" directory
+3. Download these files from this [Dropbox link](https://www.dropbox.com/sh/e5g12v7zu7sgzf7/AACW272AqPZJIUC7-A40LAsNa?dl=0) and paste them in the "resources" directory
+4. The directory structure should look like this now:
+```
+simple_qa_rnn
+  ├── datasets
+  │   └── SimpleQuestions_v2
+  │       ├── ...
+  ├── model.py
+  ├── README.md
+  ├── resources
+  │   ├── rel_to_ix_SQ.pkl
+  │   ├── w2v_map_SQ.pkl
+  │   └── word_to_ix_SQ.pkl
+  ├── saved_checkpoints
+  │   └── [...models will be saved here later...]
+  ├── scripts
+  │   ├── ...
+  ├── train.py
+  └── util.py
+```
+5. Please take a look at the arguments in utils.py and set them accordingly to train the model.
+6. Run this command to train the model. Make sure you have PyTorch and other Python dependencies installed.
+```
+python train.py 
+```
+
+NOTE: There are pre-trained models saved in the 'finished_checkpoints' directory. They can be loaded up using PyTorch.
+You can run a pre-trained model on the test dataset:
+```
+python train.py --not_bidirectional --resume_snapshot finished_checkpoints/lstm1/[model_filename] --test 
+```
@@ -0,0 +1,80 @@
+import nltk
+import string
+import pickle
+import numpy as np
+import torch
+from torch.autograd import Variable
+
+def get_all_lines(data_filename):
+    all_lines = []
+    with open(data_filename) as fin:
+        for line in fin:
+            all_lines.append(line.rstrip())
+    return all_lines
+
+def create_rp_dataset(data_file):
+    dataset = []
+    all_lines = get_all_lines(data_file)
+    for line in all_lines:
+        line_split = line.split("\t")
+        text = line_split[3]
+        relation = line_split[1]
+        dataset.append( (text, relation) )
+    return np.array(dataset)
+
+def tokenize_text(text):
+    punc_remover = str.maketrans('', '', string.punctuation)
+    processed_text = text.lower().translate(punc_remover)
+    tokens = nltk.word_tokenize(processed_text)
+    return tokens
+
+def add_padding_tokens(text_tokens, max_length, pad_type='both', pad_token='<pad>'):
+    num_pads = max_length - len(text_tokens)
+    right_pad = int(num_pads / 2)
+    left_pad = num_pads - right_pad
+    if pad_type == "both":
+        padded_tokens = [pad_token]*left_pad + text_tokens + [pad_token]*right_pad
+    elif pad_type == "right":
+        padded_tokens = text_tokens + [pad_token]*num_pads
+    else:
+        padded_tokens = [pad_token]*num_pads + text_tokens
+    return padded_tokens
+
+def load_map(pname):
+    ret_map = None
+    with open(pname, 'rb') as fh:
+        ret_map = pickle.load(fh)
+    return ret_map
+
+def text_to_vector(text, w2v_map, pad=False, max_length=None):
+    vec = []
+    tokens = tokenize_text(text)
+    if pad and (max_length != None):
+        tokens = add_padding_tokens(tokens, max_length)
+    for token in tokens:
+        vec.append( w2v_map[token] )
+    return np.array(vec)
+
+def label_to_vector(label_ix, num_labels):
+    # create one-hot vector label representation
+    y_vec = np.zeros(num_labels, dtype=np.int32)
+    y_vec[label_ix] = 1
+    return y_vec
+
+def create_tensorized_data(sentence, label, w2v_map, label_to_ix):
+    # x.shape: |S| X |D| - sentence length can vary between examples, dimension is fixed
+    x = text_to_vector(sentence, w2v_map)
+    y = label_to_ix[label]
+    inputs = Variable(torch.Tensor(x))
+    targets = Variable(torch.LongTensor([y]))
+    return inputs, targets
+
+def create_tensorized_batch(batch, max_sent_length, w2v_map, label_to_ix):
+    X = []
+    y = []
+    for sent, label in batch:
+        X.append( text_to_vector(sent, w2v_map, pad=True, max_length=max_sent_length) )
+        y.append( label_to_ix[label] )
+    inputs = Variable(torch.Tensor(X))
+    targets = Variable(torch.LongTensor(y))
+    return inputs, targets
@@ -0,0 +1,37 @@
+import torch
+import torch.nn as nn
+from torch.autograd import Variable
+import torch.nn.functional as F
+
+class BiLSTM(nn.Module):
+
+    def __init__(self, config):
+        super(BiLSTM, self).__init__()
+        self.config = config
+
+        self.lstm = nn.LSTM(input_size=config.d_embedding, hidden_size=config.d_hidden,
+                        num_layers=config.n_layers, dropout=config.dropout_prob,
+                        bidirectional=config.birnn, batch_first=True)
+
+        # linear layer maps from hidden state space to label space
+        self.hidden2label = nn.Linear(config.n_layers*config.n_directions*config.d_hidden, config.d_out)
+        self.hidden = self.init_hidden()
+        # self.dropout = nn.Dropout(p=config.dropout_prob)
+        # self.log_softmax = nn.LogSoftmax()
+
+
+    def init_hidden(self):
+        # axes semantics are (num_layers, batch_size, hidden_dim)
+        n_layers = self.config.n_layers * self.config.n_directions
+        return (Variable(torch.zeros(n_layers, self.config.batch_size, self.config.d_hidden)),
+                        Variable(torch.zeros(n_layers, self.config.batch_size, self.config.d_hidden)))
+
+    # embeds is Variable of size - (|B|, |S|, |D|)
+    def forward(self, embeds):
+        batch_size = embeds.data.size()[0]
+        sequence_length = embeds.data.size()[1]
+        lstm_out, self.hidden = self.lstm(embeds, self.hidden)
+        # print("ht size: {}".format(ht.size()))
+        rel_space = self.hidden2label(self.hidden[0].transpose(0, 1).contiguous().view(batch_size, -1)) # size - (|B|, |K|)
+        scores = F.log_softmax(rel_space)
+        return scores
@@ -0,0 +1,72 @@
+import gensim
+import glob
+import pickle
+import nltk
+import string
+import numpy as np
+
+fnames = glob.glob('../datasets/SimpleQuestions_v2/annotated*.txt')
+all_lines = []
+for fname in fnames:
+    with open(fname) as fin:
+        for line in fin:
+            all_lines.append(line.rstrip())
+
+print("num of examples in train/val/test: {}".format(len(all_lines)))
+
+max_sent_length = -1
+for line in all_lines:
+    qText = line.split("\t")[3]
+    punc_remover = str.maketrans('', '', string.punctuation)
+    processed_text = qText.lower().translate(punc_remover)
+    tokens = nltk.word_tokenize(processed_text)
+    if len(tokens) > max_sent_length:
+        max_sent_length = len(tokens)
+
+print("max_sent_length: {}".format(max_sent_length))
+
+
+
+# # get all the words in the train, dev, val, test set
+# all_words = set()
+# for line in all_lines:
+#     qText = line.split("\t")[3]
+#     # process text: remove punctuations, lowercase
+#     punc_remover = str.maketrans('', '', string.punctuation)
+#     processed_text = qText.lower().translate(punc_remover)
+#     tokens = nltk.word_tokenize(processed_text)
+#     for tok in tokens:
+#         all_words.add(tok)
+#
+# print("vocab. size: {}".format(len(all_words)))
+#
+# # word_to_ix = { word:i for i, word in enumerate(all_words) }
+# # # dump the pickle
+# # with open("../resources/word_to_ix_SQ.pkl", 'wb') as fh:
+# #     pickle.dump(word_to_ix, fh)
+#
+# word_to_vector_map = {}
+# w2v_path = "../resources/GoogleNews-vectors-negative300.bin.gz"
+# # get their word vectors
+# print("loading word vectors...")
+# word_vectors = gensim.models.KeyedVectors.load_word2vec_format(w2v_path, binary=True)
+#
+# # store in a dict
+# found = 0
+# random = 0
+# for w in all_words:
+#     # print(w)
+#     try:
+#         word_to_vector_map[w] = word_vectors[w]
+#         found += 1
+#     except:
+#         word_to_vector_map[w] = np.random.uniform(low=0.0, high=1.0, size=300)
+#         random += 1
+#
+# print("found: {}".format(found))
+# print("random: {}".format(random))
+#
+# # dump the pickle
+# print("dumping the w2v map pickle...")
+# with open("../resources/w2v_map_SQ.pkl", 'wb') as fh:
+#     pickle.dump(word_to_vector_map, fh)
@@ -0,0 +1,56 @@
+import gensim
+import glob
+import json
+import pickle
+import nltk
+import string
+import numpy as np
+
+fnames = glob.glob('../datasets/dataset-factoid-webquestions/main/*.json')
+all_entries = []
+for fname in fnames:
+    with open(fname) as fin:
+        data = json.load(fin)
+        all_entries.extend(data)
+
+print("num of examples in train/val/test: {}".format(len(all_entries)))
+
+# get all the words in the train, dev, val, test set
+all_words = set()
+for entry in all_entries:
+    qText = entry.get('qText')
+    # process text: remove punctuations, lowercase
+    punc_remover = str.maketrans('', '', string.punctuation)
+    processed_text = qText.lower().translate(punc_remover)
+    tokens = nltk.word_tokenize(processed_text)
+    for tok in tokens:
+        all_words.add(tok)
+
+print("vocab. size: {}".format(len(all_words)))
+
+word_to_vector_map = {}
+
+w2v_path = "../resources/GoogleNews-vectors-negative300.bin.gz"
+# get their word vectors
+print("loading word vectors...")
+word_vectors = gensim.models.KeyedVectors.load_word2vec_format(w2v_path, binary=True)
+
+# store in a dict
+found = 0
+random = 0
+for w in all_words:
+    # print(w)
+    try:
+        word_to_vector_map[w] = word_vectors[w]
+        found += 1
+    except:
+        word_to_vector_map[w] = np.random.uniform(low=0.0, high=1.0, size=300)
+        random += 1
+
+print("found: {}".format(found))
+print("random: {}".format(random))
+
+# dump the pickle
+print("dumping the w2v map pickle...")
+with open("../resources/w2v_map_WQ.pkl", 'wb') as fh:
+    pickle.dump(word_to_vector_map, fh)
@@ -0,0 +1,11 @@
+import pickle
+
+# load the pickled word vectors
+w2v_pkl_path = "w2v_map_SQ.pkl"
+print("loading word vectors from the pickle...")
+w2v_map = None
+with open(w2v_pkl_path, 'rb') as fh:
+    w2v_map = pickle.load(fh)
+
+print(list(w2v_map.keys()))
+print(len(w2v_map.keys()))
@@ -0,0 +1,27 @@
+import glob
+import pickle
+
+fnames = glob.glob('../datasets/SimpleQuestions_v2/annotated*.txt')
+all_lines = []
+for fname in fnames:
+    with open(fname) as fin:
+        for line in fin:
+            all_lines.append(line.rstrip())
+
+print("num of examples in train/val/test: {}".format(len(all_lines)))
+
+all_relations = set()
+for line in all_lines:
+    rel = line.split("\t")[1]
+    all_relations.add(rel)
+
+print("num of relation types: {}".format(len(all_relations)))
+print(all_relations)
+print( ("www.freebase.com/music/release/region") in all_relations )
+
+rel_to_ix = { rel:i for i, rel in enumerate(all_relations) }
+# print(rel_to_ix)
+# # dump the pickle
+# print("dumping the w2v map pickle...")
+# with open("../resources/rel_to_ix_SQ.pkl", 'wb') as fh:
+#     pickle.dump(rel_to_ix, fh)
@@ -0,0 +1,24 @@
+import glob
+import json
+
+fnames = glob.glob('dataset-factoid-webquestions/d-freebase-rp/*.json')
+all_entries = []
+for fname in fnames:
+    with open(fname) as fin:
+        data = json.load(fin)
+        all_entries.extend(data)
+
+print("num of examples in train/val/test: {}".format(len(all_entries)))
+
+all_relations = set()
+for entry in all_entries:
+    qID = entry.get('qId')
+    relPaths = entry.get('relPaths')
+    for relPath in relPaths:
+        relations = relPath[0]
+        for rel in relations:
+            all_relations.add(rel)
+
+print("num of relation types: {}".format(len(all_relations)))
+print(all_relations)
+print( ("/people/marriage/spouse") in all_relations )
@@ -0,0 +1,137 @@
+import os
+import sys
+import time
+import glob
+import numpy as np
+import torch
+from torch.autograd import Variable
+import torch.nn as nn
+import torch.optim as optim
+from model import BiLSTM
+from util import get_args
+import data
+
+args = get_args()
+
+
+# torch.cuda.set_device(args.gpu)
+
+
+# ---- Helper Methods ------
+def evaluate_dataset_batch(data_set, max_sent_length, model, w2v_map, label_to_ix):
+    n_total = len(data_set)
+    n_correct = 0
+    num_batches = len(data_set) // args.batch_size
+    batch_indices = np.split(range(n_total),
+                                range(args.batch_size, n_total, args.batch_size))
+    model.eval()
+    for batch_ix in range(num_batches):
+        batch = data_set[batch_indices[batch_ix]]
+        inputs, targets = data.create_tensorized_batch(batch, max_sent_length, w2v_map, label_to_ix)
+        scores = model(inputs)
+        pred_label_ix = np.argmax(scores.data.numpy(), axis=1) # check this properly
+        correct_label_ix = targets.data.numpy()
+        n_correct += (pred_label_ix == correct_label_ix).sum()
+    acc = n_correct / n_total
+    return acc
+
+def repackage_hidden(h):
+    """Wraps hidden states in new Variables, to detach them from their history."""
+    if type(h) == Variable:
+        return Variable(h.data)
+    else:
+        return tuple(repackage_hidden(v) for v in h)
+
+#Load Datasets ------
+train_file = "datasets/SimpleQuestions_v2/annotated_fb_data_train.txt"
+val_file = "datasets/SimpleQuestions_v2/annotated_fb_data_valid.txt"
+test_file = "datasets/SimpleQuestions_v2/annotated_fb_data_test.txt"
+
+train_set = data.create_rp_dataset(train_file)
+val_set = data.create_rp_dataset(val_file)
+test_set = data.create_rp_dataset(test_file)
+# train_set = train_set[:4]  # work with few examples first
+
+# ---- Build Vocabulary ------
+w2v_map = data.load_map("resources/w2v_map_SQ.pkl")
+w2v_map['<pad>'] = np.zeros(300)
+word_to_ix = data.load_map("resources/word_to_ix_SQ.pkl")
+label_to_ix = data.load_map("resources/rel_to_ix_SQ.pkl")
+vocab_size = len(word_to_ix)
+num_classes = len(label_to_ix)
+max_sent_length = 36  # set from the paper
+
+# ---- Define Model, Loss, Optim ------
+config = args
+config.d_out = num_classes
+config.n_directions = 2 if config.birnn else 1
+print(config)
+model = BiLSTM(config)
+loss_function = nn.NLLLoss()
+optimizer = optim.Adam(model.parameters(), lr=args.lr)
+
+# ---- Test Model ------
+if args.test:
+    print("Test Mode: loading pre-trained model and testing on test set...")
+    # model = torch.load(args.resume_snapshot, map_location=lambda storage, location: storage.cuda(args.gpu))
+    model.load_state_dict(torch.load(args.resume_snapshot))
+    test_acc = evaluate_dataset_batch(test_set, max_sent_length, model, w2v_map, label_to_ix)
+    print("Accuracy: {}".format(test_acc))
+    sys.exit(0)
+
+
+# ---- Train Model ------
+start = time.time()
+best_val_acc = -1
+iter = 0
+header = '  Time Epoch Iteration     Loss   Train/Acc.   Val/Acc.'
+print(header)
+log_template = ' '.join('{:>6.0f},{:>5.0f},{:>9.0f},{:>9.6f}'.split(','))
+dev_log_template = ' '.join('{:>6.0f},{:>5.0f},{:>9.0f},{:>9.6f},{:9.6f},{:11.6f}'.split(','))
+
+model.train()
+for epoch in range(args.epochs):
+    # shuffle the dataset and create batches (truncate the last batch if not of equal size)
+    shuffled_indices = np.random.permutation(len(train_set))
+    num_batches = len(shuffled_indices) // args.batch_size
+    batch_indices = np.split(shuffled_indices,
+                                range(args.batch_size, len(shuffled_indices), args.batch_size))
+    model.hidden = model.init_hidden()
+    for batch_ix in range(num_batches):
+        iter += 1
+        batch = train_set[batch_indices[batch_ix]]
+        inputs, targets = data.create_tensorized_batch(batch, max_sent_length, w2v_map, label_to_ix)
+        # print("inputs size: {}".format(inputs.size()))
+        # print("targets size: {}".format(targets.size()))
+
+        # clear out gradients and hidden states of the model
+        model.zero_grad()
+        model.hidden = repackage_hidden(model.hidden)
+
+        # prepare inputs for LSTM model and run forward pass
+        scores = model(inputs)
+
+        # compute the loss, gradients, and update the parameters
+        loss = loss_function(scores, targets)
+        loss.backward()
+        # `clip_grad_norm` helps prevent the exploding gradient problem in RNNs / LSTMs.
+        torch.nn.utils.clip_grad_norm(model.parameters(), args.clip)
+        optimizer.step()
+
+        # log at intervals
+        if iter % args.dev_every == 0:
+            train_acc = evaluate_dataset_batch(train_set[:8000], max_sent_length, model, w2v_map, label_to_ix)
+            val_acc = evaluate_dataset_batch(val_set, max_sent_length, model, w2v_map, label_to_ix)
+            print(dev_log_template.format(time.time()-start, epoch, iter, loss.data[0], train_acc, val_acc))
+            if val_acc > best_val_acc:
+                best_val_acc = val_acc
+                snapshot_prefix = os.path.join(args.save_path, 'best_snapshot')
+                snapshot_path = snapshot_prefix + '_valacc_{:6.4f}__iter_{}_model.pt'.format(val_acc, iter)
+                torch.save(model.state_dict(), snapshot_path)
+                for f in glob.glob(snapshot_prefix + '*'):
+                    if f != snapshot_path:
+                        os.remove(f)
+
+        elif iter == 1 or iter % args.log_every == 0:
+            print(log_template.format(time.time() - start, epoch, iter, loss.data[0]))
+
@@ -0,0 +1,22 @@
+from argparse import ArgumentParser
+
+def get_args():
+    parser = ArgumentParser(description='Simple QA BiGRU model - Ferhan')
+    parser.add_argument('--epochs', type=int, default=30)
+    parser.add_argument('--batch_size', type=int, default=256)
+    parser.add_argument('--d_embedding', type=int, default=300)
+    parser.add_argument('--d_hidden', type=int, default=300)
+    parser.add_argument('--n_layers', type=int, default=1)
+    parser.add_argument('--test', action='store_true', dest='test', help='turn on test mode; no training.')
+    parser.add_argument('--not_bidirectional', action='store_false', dest='birnn')
+    parser.add_argument('--clip', type=float, default=0.25, help='gradient clipping')
+    parser.add_argument('--log_every', type=int, default=50)
+    parser.add_argument('--lr', type=float, default=1e-3)
+    parser.add_argument('--dev_every', type=int, default=300)
+    parser.add_argument('--save_every', type=int, default=1000)
+    parser.add_argument('--dropout_prob', type=int, default=0.2)
+    parser.add_argument('--gpu', type=int, default=0)
+    parser.add_argument('--save_path', type=str, default='saved_checkpoints')
+    parser.add_argument('--resume_snapshot', type=str, default='')
+    args = parser.parse_args()
+    return args