snli_rnn.py

# Adopted code from https://github.com/Smerity/keras_snli
# Modified by Moustafa Alzantot (malzantot@ucla.edu)

from __future__ import print_function
from keras.utils import np_utils
from keras.regularizers import l2
from keras.preprocessing.text import Tokenizer
from keras.preprocessing.sequence import pad_sequences
from keras.models import Model
from keras.layers.wrappers import Bidirectional
from keras.layers.normalization import BatchNormalization
from keras.layers.embeddings import Embedding
from keras.layers import merge, recurrent, Dense, Input, Dropout, TimeDistributed
from keras.callbacks import EarlyStopping, ModelCheckpoint
import keras.backend as K
import keras
from functools import reduce
import json
import os
import re
import tarfile
import tempfile
import pickle

import numpy as np
np.random.seed(1337)  # for reproducibility

'''
300D Model - Train / Test (epochs)
=-=-=
Batch size = 512
Fixed GloVe
- 300D SumRNN + Translate + 3 MLP (1.2 million parameters) - 0.8315 / 0.8235 / 0.8249 (22 epochs)
- 300D GRU + Translate + 3 MLP (1.7 million parameters) - 0.8431 / 0.8303 / 0.8233 (17 epochs)
- 300D LSTM + Translate + 3 MLP (1.9 million parameters) - 0.8551 / 0.8286 / 0.8229 (23 epochs)

Following Liu et al. 2016, I don't update the GloVe embeddings during training.
Unlike Liu et al. 2016, I don't initialize out of vocabulary embeddings randomly and instead leave them zeroed.

The jokingly named SumRNN (summation of word embeddings) is 10-11x faster than the GRU or LSTM.

Original numbers for sum / LSTM from Bowman et al. '15 and Bowman et al. '16
=-=-=
100D Sum + GloVe - 0.793 / 0.753
100D LSTM + GloVe - 0.848 / 0.776
300D LSTM + GloVe - 0.839 / 0.806
'''


def extract_tokens_from_binary_parse(parse):
    return parse.replace('(', ' ').replace(')', ' ').replace('-LRB-', '(').replace('-RRB-', ')').split()


def yield_examples(fn, skip_no_majority=True, limit=None):
    for i, line in enumerate(open(fn)):
        if limit and i > limit:
            break
        data = json.loads(line)
        label = data['gold_label']
        s1 = ' '.join(extract_tokens_from_binary_parse(
            data['sentence1_binary_parse']))
        s2 = ' '.join(extract_tokens_from_binary_parse(
            data['sentence2_binary_parse']))
        if skip_no_majority and label == '-':
            continue
        yield (label, s1, s2)


def get_data(fn, limit=None):
    raw_data = list(yield_examples(fn=fn, limit=limit))
    left = [s1 for _, s1, s2 in raw_data]
    right = [s2 for _, s1, s2 in raw_data]
    print(max(len(x.split()) for x in left))
    print(max(len(x.split()) for x in right))

    LABELS = {'contradiction': 0, 'neutral': 1, 'entailment': 2}
    Y = np.array([LABELS[l] for l, s1, s2 in raw_data])
    Y = np_utils.to_categorical(Y, len(LABELS))

    return left, right, Y


training = get_data('snli_data/snli_1.0/snli_1.0_train.jsonl')
validation = get_data('snli_data/snli_1.0/snli_1.0_dev.jsonl')
test = get_data('snli_data/snli_1.0/snli_1.0_test.jsonl')

tokenizer = Tokenizer(lower=False, filters='')
tokenizer.fit_on_texts(training[0] + training[1])

# Lowest index from the tokenizer is 1 - we need to include 0 in our vocab count
VOCAB = len(tokenizer.word_counts) + 1
LABELS = {'contradiction': 0, 'neutral': 1, 'entailment': 2}
#RNN = recurrent.LSTM
#RNN = lambda *args, **kwargs: Bidirectional(recurrent.LSTM(*args, **kwargs))
#RNN = recurrent.GRU
#RNN = lambda *args, **kwargs: Bidirectional(recurrent.GRU(*args, **kwargs))
# Summation of word embeddings
RNN = None
LAYERS = 1
USE_GLOVE = True
TRAIN_EMBED = False
EMBED_HIDDEN_SIZE = 300
SENT_HIDDEN_SIZE = 300
BATCH_SIZE = 512
PATIENCE = 4  # 8
MAX_EPOCHS = 42
MAX_LEN = 42
DP = 0.2
L2 = 4e-6
ACTIVATION = 'relu'
OPTIMIZER = 'rmsprop'
print('RNN / Embed / Sent = {}, {}, {}'.format(RNN,
                                               EMBED_HIDDEN_SIZE, SENT_HIDDEN_SIZE))
print('GloVe / Trainable Word Embeddings = {}, {}'.format(USE_GLOVE, TRAIN_EMBED))


def to_seq(X): return pad_sequences(
    tokenizer.texts_to_sequences(X), maxlen=MAX_LEN)


def prepare_data(data): return (to_seq(data[0]), to_seq(data[1]), data[2])


def build_model():

    GLOVE_STORE = 'nli_precomputed_glove.weights'
    if USE_GLOVE:
        if not os.path.exists(GLOVE_STORE + '.npy'):
            print('Computing GloVe')
            embeddings_index = {}
            f = open('glove.840B.300d.txt')
            for line in f:
                values = line.split(' ')
                word = values[0]
                coefs = np.asarray(values[1:], dtype='float32')
                embeddings_index[word] = coefs
            f.close()
            # prepare embedding matrix
            embedding_matrix = np.zeros((VOCAB, EMBED_HIDDEN_SIZE))
            for word, i in tokenizer.word_index.items():
                embedding_vector = embeddings_index.get(word)
                if embedding_vector is not None:
                    # words not found in embedding index will be all-zeros.
                    embedding_matrix[i] = embedding_vector
                else:
                    print('Missing from GloVe: {}'.format(word))
            np.save(GLOVE_STORE, embedding_matrix)
        print('Loading GloVe')
        embedding_matrix = np.load(GLOVE_STORE + '.npy')
        print('Total number of null word embeddings:')
        print(np.sum(np.sum(embedding_matrix, axis=1) == 0))
        embed = Embedding(VOCAB, EMBED_HIDDEN_SIZE, weights=[
            embedding_matrix], input_length=MAX_LEN, trainable=TRAIN_EMBED)
    else:
        embed = Embedding(VOCAB, EMBED_HIDDEN_SIZE, input_length=MAX_LEN)
    rnn_kwargs = dict(output_dim=SENT_HIDDEN_SIZE, dropout_W=DP, dropout_U=DP)
    SumEmbeddings = keras.layers.core.Lambda(
        lambda x: K.sum(x, axis=1), output_shape=(SENT_HIDDEN_SIZE, ))
    translate = TimeDistributed(Dense(SENT_HIDDEN_SIZE, activation=ACTIVATION))

    premise = Input(shape=(MAX_LEN,), dtype='int32')
    hypothesis = Input(shape=(MAX_LEN,), dtype='int32')
    prem = embed(premise)
    hypo = embed(hypothesis)
    prem = translate(prem)
    hypo = translate(hypo)
    if RNN and LAYERS > 1:
        for l in range(LAYERS - 1):
            rnn = RNN(return_sequences=True, **rnn_kwargs)
            prem = BatchNormalization()(rnn(prem))
            hypo = BatchNormalization()(rnn(hypo))
    rnn = SumEmbeddings if not RNN else RNN(
        return_sequences=False, **rnn_kwargs)
    prem = rnn(prem)
    hypo = rnn(hypo)
    prem = BatchNormalization()(prem)
    hypo = BatchNormalization()(hypo)
    joint = merge.Concatenate()([prem, hypo])
    joint = Dropout(DP)(joint)
    for i in range(3):
        joint = Dense(2 * SENT_HIDDEN_SIZE, activation=ACTIVATION,
                      W_regularizer=l2(L2) if L2 else None)(joint)
        joint = Dropout(DP)(joint)
        joint = BatchNormalization()(joint)

    pred = Dense(len(LABELS), activation='softmax')(joint)
    model = Model(input=[premise, hypothesis], output=pred)
    model.compile(optimizer=OPTIMIZER,
                  loss='categorical_crossentropy', metrics=['accuracy'])

    return model


if __name__ == '__main__':
    training = prepare_data(training)
    validation = prepare_data(validation)
    test = prepare_data(test)
    print('Vocab size =', VOCAB)
    model = build_model()
    print('Build model...')
    model.summary()
    print('Training')
    _, tmpfn = tempfile.mkstemp()
    # Save the best model during validation and bail out of training early if we're not improving
    callbacks = [EarlyStopping(patience=PATIENCE), ModelCheckpoint(
        tmpfn, save_best_only=True, save_weights_only=True)]
    model.fit([training[0], training[1]], training[2], batch_size=BATCH_SIZE, nb_epoch=MAX_EPOCHS,
              validation_data=([validation[0], validation[1]], validation[2]), callbacks=callbacks)
    # Restore the best found model during validation
    model.load_weights(tmpfn)
    loss, acc = model.evaluate(
        [test[0], test[1]], test[2], batch_size=BATCH_SIZE)
    print('Test loss / test accuracy = {:.4f} / {:.4f}'.format(loss, acc))
    # Save the model
    model.save_weights('nli_model.h5')
    with open('nli_tokenizer.pkl', 'wb') as fh:
        pickle.dump(tokenizer, fh)
    with open('nli_testing.pkl', 'wb') as fh:
        pickle.dump(test, fh)