diff --git a/attention_lstm.py b/attention_lstm.py
index 03de71c..f198942 100644
--- a/attention_lstm.py
+++ b/attention_lstm.py
@@ -1,16 +1,15 @@
 from __future__ import absolute_import
 
 from keras import backend as K
-from keras.layers import LSTM, activations
+from keras.engine import InputSpec
+from keras.layers import LSTM, activations, Wrapper
 
 
 class AttentionLSTM(LSTM):
-    def __init__(self, output_dim, attention_vec, attn_activation='tanh',
-                 single_attention_param=False, n_attention_dim=None, **kwargs):
+    def __init__(self, output_dim, attention_vec, attn_activation='tanh', single_attention_param=False, **kwargs):
         self.attention_vec = attention_vec
         self.attn_activation = activations.get(attn_activation)
         self.single_attention_param = single_attention_param
-        self.n_attention_dim = output_dim if n_attention_dim is None else n_attention_dim
 
         super(AttentionLSTM, self).__init__(output_dim, **kwargs)
 
@@ -65,3 +64,107 @@ class AttentionLSTM(LSTM):
         constants = super(AttentionLSTM, self).get_constants(x)
         constants.append(K.dot(self.attention_vec, self.U_m) + self.b_m)
         return constants
+
+
+class AttentionLSTMWrapper(Wrapper):
+    def __init__(self, layer, attention_vec, attn_activation='tanh', single_attention_param=False, **kwargs):
+        assert isinstance(layer, LSTM)
+        self.supports_masking = True
+        self.attention_vec = attention_vec
+        self.attn_activation = activations.get(attn_activation)
+        self.single_attention_param = single_attention_param
+        super(AttentionLSTMWrapper, self).__init__(layer, **kwargs)
+
+    def build(self, input_shape):
+        assert len(input_shape) >= 3
+        self.input_spec = [InputSpec(shape=input_shape)]
+
+        if not self.layer.built:
+            self.layer.build(input_shape)
+            self.layer.built = True
+
+        super(AttentionLSTMWrapper, self).build()
+
+        if hasattr(self.attention_vec, '_keras_shape'):
+            attention_dim = self.attention_vec._keras_shape[1]
+        else:
+            raise Exception('Layer could not be build: No information about expected input shape.')
+
+        self.U_a = self.layer.inner_init((self.layer.output_dim, self.layer.output_dim), name='{}_U_a'.format(self.name))
+        self.b_a = K.zeros((self.layer.output_dim,), name='{}_b_a'.format(self.name))
+
+        self.U_m = self.layer.inner_init((attention_dim, self.layer.output_dim), name='{}_U_m'.format(self.name))
+        self.b_m = K.zeros((self.layer.output_dim,), name='{}_b_m'.format(self.name))
+
+        if self.single_attention_param:
+            self.U_s = self.layer.inner_init((self.layer.output_dim, 1), name='{}_U_s'.format(self.name))
+            self.b_s = K.zeros((1,), name='{}_b_s'.format(self.name))
+        else:
+            self.U_s = self.layer.inner_init((self.layer.output_dim, self.layer.output_dim), name='{}_U_s'.format(self.name))
+            self.b_s = K.zeros((self.layer.output_dim,), name='{}_b_s'.format(self.name))
+
+        self.trainable_weights = [self.U_a, self.U_m, self.U_s, self.b_a, self.b_m, self.b_s]
+
+    def get_output_shape_for(self, input_shape):
+        return self.layer.get_output_shape_for(input_shape)
+
+    def step(self, x, states):
+        h, [h, c] = self.layer.step(x, states)
+        attention = states[4]
+
+        m = self.attn_activation(K.dot(h, self.U_a) * attention + self.b_a)
+        s = K.sigmoid(K.dot(m, self.U_s) + self.b_s)
+
+        if self.single_attention_param:
+            h = h * K.repeat_elements(s, self.layer.output_dim, axis=1)
+        else:
+            h = h * s
+
+        return h, [h, c]
+
+    def get_constants(self, x):
+        constants = self.layer.get_constants(x)
+        constants.append(K.dot(self.attention_vec, self.U_m) + self.b_m)
+        return constants
+
+    def call(self, x, mask=None):
+        # input shape: (nb_samples, time (padded with zeros), input_dim)
+        # note that the .build() method of subclasses MUST define
+        # self.input_spec with a complete input shape.
+        input_shape = self.input_spec[0].shape
+        if K._BACKEND == 'tensorflow':
+            if not input_shape[1]:
+                raise Exception('When using TensorFlow, you should define '
+                                'explicitly the number of timesteps of '
+                                'your sequences.\n'
+                                'If your first layer is an Embedding, '
+                                'make sure to pass it an "input_length" '
+                                'argument. Otherwise, make sure '
+                                'the first layer has '
+                                'an "input_shape" or "batch_input_shape" '
+                                'argument, including the time axis. '
+                                'Found input shape at layer ' + self.name +
+                                ': ' + str(input_shape))
+        if self.layer.stateful:
+            initial_states = self.layer.states
+        else:
+            initial_states = self.layer.get_initial_states(x)
+        constants = self.get_constants(x)
+        preprocessed_input = self.layer.preprocess_input(x)
+
+        last_output, outputs, states = K.rnn(self.step, preprocessed_input,
+                                             initial_states,
+                                             go_backwards=self.layer.go_backwards,
+                                             mask=mask,
+                                             constants=constants,
+                                             unroll=self.layer.unroll,
+                                             input_length=input_shape[1])
+        if self.layer.stateful:
+            self.updates = []
+            for i in range(len(states)):
+                self.updates.append((self.layer.states[i], states[i]))
+
+        if self.layer.return_sequences:
+            return outputs
+        else:
+            return last_output
diff --git a/insurance_qa_embeddings.py b/insurance_qa_embeddings.py
deleted file mode 100644
index f6a18b5..0000000
--- a/insurance_qa_embeddings.py
+++ /dev/null
@@ -1,56 +0,0 @@
-from __future__ import print_function
-
-import os
-import sys
-import random
-
-import pickle
-
-from gensim.models import Word2Vec
-
-from keras_models import *
-
-random.seed(42)
-
-
-def load(path, name):
-    return pickle.load(open(os.path.join(path, name), 'rb'))
-
-
-def revert(vocab, indices):
-    return [vocab.get(i, 'X') for i in indices]
-
-if __name__ == '__main__':
-    try:
-        data_path = os.environ['INSURANCE_QA']
-    except KeyError:
-        print("INSURANCE_QA is not set.  Set it to your clone of https://github.com/codekansas/insurance_qa_python")
-        sys.exit(1)
-
-    size = 1000
-    assert os.path.exists('models/embedding_%d_dim.h5' % size)
-        
-    vocab = load(data_path, 'vocabulary')
-
-    sentences = list()
-    answers = load(data_path, 'answers')
-    for id, txt in answers.items():
-        sentences.append(revert(vocab, txt))
-    for q in load(data_path, 'train'):
-        sentences.append(revert(vocab, q['question']))
-
-    print('Training Word2Vec model...')
-    model = Word2Vec(sentences, size=size, min_count=5, window=5, sg=1, iter=25)
-    weights = model.syn0
-    d = dict([(k, v.index) for k, v in model.vocab.items()])
-
-    # this is the stored weights of an equivalent embedding layer
-    # there is some commented code in insurance_qa_eval.py for generating this
-    emb = np.load('models/embedding_%d_dim.h5' % size)
-
-    # swap the word2vec weights with the embedded weights
-    for i, w in vocab.items():
-        if w not in d: continue
-        emb[i, :] = weights[d[w], :]
-
-    np.save(open('models/word2vec_%d_dim.h5' % size, 'wb'), emb)
diff --git a/insurance_qa_eval.py b/insurance_qa_eval.py
index a822af0..b4656a1 100644
--- a/insurance_qa_eval.py
+++ b/insurance_qa_eval.py
@@ -7,25 +7,29 @@ import random
 from time import strftime, gmtime
 
 import pickle
+import json
 
-from keras.optimizers import Adam
+from keras.optimizers import SGD
 from scipy.stats import rankdata
 
-from keras_models import EmbeddingModel, AttentionModel, ConvolutionModel
-
 random.seed(42)
 
 
 class Evaluator:
-    def __init__(self, conf=None):
+    def __init__(self, conf, model=None, optimizer=None):
         try:
             data_path = os.environ['INSURANCE_QA']
         except KeyError:
-            print("INSURANCE_QA is not set.  Set it to your clone of https://github.com/codekansas/insurance_qa_python")
+            print("INSURANCE_QA is not set. Set it to your clone of https://github.com/codekansas/insurance_qa_python")
             sys.exit(1)
+        if isinstance(conf, str):
+            conf = json.load(open(conf, 'rb'))
+        self.model = conf['model'](conf) if model is None else model
         self.path = data_path
-        self.conf = dict() if conf is None else conf
-        self.params = conf.get('training_params', dict())
+        self.conf = conf
+        self.params = conf['training']
+        optimizer = self.params['optimizer'] if optimizer is None else optimizer
+        self.model.compile(optimizer)
         self.answers = self.load('answers') # self.load('generated')
         self._vocab = None
         self._reverse_vocab = None
@@ -49,14 +53,14 @@ class Evaluator:
 
     ##### Loading / saving #####
 
-    def save_epoch(self, model, epoch):
+    def save_epoch(self, epoch):
         if not os.path.exists('models/'):
             os.makedirs('models/')
-        model.save_weights('models/weights_epoch_%d.h5' % epoch, overwrite=True)
+        self.model.save_weights('models/weights_epoch_%d.h5' % epoch, overwrite=True)
 
-    def load_epoch(self, model, epoch):
+    def load_epoch(self, epoch):
         assert os.path.exists('models/weights_epoch_%d.h5' % epoch), 'Weights at epoch %d not found' % epoch
-        model.load_weights('models/weights_epoch_%d.h5' % epoch)
+        self.model.load_weights('models/weights_epoch_%d.h5' % epoch)
 
     ##### Converting / reverting #####
 
@@ -87,41 +91,49 @@ class Evaluator:
     def print_time(self):
         print(strftime('%Y-%m-%d %H:%M:%S :: ', gmtime()), end='')
 
-    def train(self, model):
-        save_every = self.params.get('save_every', None)
-        batch_size = self.params.get('batch_size', 128)
-        nb_epoch = self.params.get('nb_epoch', 10)
-        split = self.params.get('validation_split', 0)
+    def train(self):
+        batch_size = self.params['batch_size']
+        nb_epoch = self.params['nb_epoch']
+        validation_split = self.params['validation_split']
 
         training_set = self.load('train')
+        top_50 = self.load('top_50')
 
         questions = list()
         good_answers = list()
+        indices = list()
 
-        for q in training_set:
+        for j, q in enumerate(training_set):
             questions += [q['question']] * len(q['answers'])
             good_answers += [self.answers[i] for i in q['answers']]
+            indices += [j] * len(q['answers'])
 
         questions = self.padq(questions)
         good_answers = self.pada(good_answers)
 
         val_loss = {'loss': 1., 'epoch': 0}
 
+        def get_bad_samples(indices, top_50):
+            return [self.answers[random.choice(top_50[i])] for i in indices]
+
         for i in range(1, nb_epoch):
             # sample from all answers to get bad answers
+            # if i % 2 == 0:
+            #     bad_answers = self.pada(random.sample(self.answers.values(), len(good_answers)))
+            # else:
+            #     bad_answers = self.pada(get_bad_samples(indices, top_50))
             bad_answers = self.pada(random.sample(self.answers.values(), len(good_answers)))
 
             print('Epoch %d :: ' % i, end='')
             self.print_time()
-            hist = model.fit([questions, good_answers, bad_answers], nb_epoch=1, batch_size=batch_size,
-                             validation_split=split)
+            hist = self.model.fit([questions, good_answers, bad_answers], nb_epoch=1, batch_size=batch_size,
+                             validation_split=validation_split)
 
             if hist.history['val_loss'][0] < val_loss['loss']:
                 val_loss = {'loss': hist.history['val_loss'][0], 'epoch': i}
             print('Best: Loss = {}, Epoch = {}'.format(val_loss['loss'], val_loss['epoch']))
 
-            if save_every is not None and i % save_every == 0:
-                self.save_epoch(model, i)
+            self.save_epoch(i)
 
         return val_loss
 
@@ -140,40 +152,43 @@ class Evaluator:
             self._eval_sets = dict([(s, self.load(s)) for s in ['dev', 'test1', 'test2']])
         return self._eval_sets
 
-    def get_mrr(self, model, evaluate_all=False):
-        top1s = list()
-        mrrs = list()
-
+    def get_score(self, verbose=False):
         for name, data in self.eval_sets().items():
-            if evaluate_all:
-                self.print_time()
-                print('----- %s -----' % name)
+            self.print_time()
+            print('----- %s -----' % name)
 
             random.shuffle(data)
 
-            if not evaluate_all and 'n_eval' in self.params:
+            if 'n_eval' in self.params:
                 data = data[:self.params['n_eval']]
 
             c_1, c_2 = 0, 0
 
             for i, d in enumerate(data):
-                if evaluate_all:
-                    self.prog_bar(i, len(data))
+                self.prog_bar(i, len(data))
 
                 indices = d['good'] + d['bad']
                 answers = self.pada([self.answers[i] for i in indices])
                 question = self.padq([d['question']] * len(indices))
 
-                n_good = len(d['good'])
-                sims = model.predict([question, answers])
-                r = rankdata(sims, method='max')
+                sims = self.model.predict([question, answers])
 
+                n_good = len(d['good'])
                 max_r = np.argmax(sims)
                 max_n = np.argmax(sims[:n_good])
 
-                # print(' '.join(self.revert(d['question'])))
-                # print(' '.join(self.revert(self.answers[indices[max_r]])))
-                # print(' '.join(self.revert(self.answers[indices[max_n]])))
+                r = rankdata(sims, method='max')
+
+                if verbose:
+                    min_r = np.argmin(sims)
+                    amin_r = self.answers[indices[min_r]]
+                    amax_r = self.answers[indices[max_r]]
+                    amax_n = self.answers[indices[max_n]]
+
+                    print(' '.join(self.revert(d['question'])))
+                    print('Predicted: ({}) '.format(sims[max_r]) + ' '.join(self.revert(amax_r)))
+                    print('Expected: ({}) Rank = {} '.format(sims[max_n], r[max_n]) + ' '.join(self.revert(amax_n)))
+                    print('Worst: ({})'.format(sims[min_r]) + ' '.join(self.revert(amin_r)))
 
                 c_1 += 1 if max_r == max_n else 0
                 c_2 += 1 / float(r[max_r] - r[max_n] + 1)
@@ -182,69 +197,27 @@ class Evaluator:
             mrr = c_2 / float(len(data))
 
             del data
-
-            if evaluate_all:
-                print('Top-1 Precision: %f' % top1)
-                print('MRR: %f' % mrr)
-
-            top1s.append(top1)
-            mrrs.append(mrr)
-
-        # rerun the evaluation if above some threshold
-        if not evaluate_all:
-            print('Top-1 Precision: {}'.format(top1s))
-            print('MRR: {}'.format(mrrs))
-            evaluate_all_threshold = self.params.get('evaluate_all_threshold', dict())
-            evaluate_mode = evaluate_all_threshold.get('mode', 'all')
-            mrr_theshold = evaluate_all_threshold.get('mrr', 1)
-            top1_threshold = evaluate_all_threshold.get('top1', 1)
-
-            if evaluate_mode == 'any':
-                evaluate_all = evaluate_all or any([x >= top1_threshold for x in top1s])
-                evaluate_all = evaluate_all or any([x >= mrr_theshold for x in mrrs])
-            else:
-                evaluate_all = evaluate_all or all([x >= top1_threshold for x in top1s])
-                evaluate_all = evaluate_all or all([x >= mrr_theshold for x in mrrs])
-
-            if evaluate_all:
-                return self.get_mrr(model, evaluate_all=True)
-
-        return top1s, mrrs
+            print('Top-1 Precision: %f' % top1)
+            print('MRR: %f' % mrr)
 
 
 if __name__ == '__main__':
     import numpy as np
 
     conf = {
-        'question_len': 50,
-        'answer_len': 100,
-        'n_words': 22353,  # len(vocabulary) + 1
-        'margin': 0.02,
+        'n_words': 22353,
+        'question_len': 150,
+        'answer_len': 150,
+        'margin': 0.05,
+        'initial_embed_weights': 'word2vec_100_dim.embeddings',
 
-        'training_params': {
-            'save_every': 1,
-            'batch_size': 20,
-            'nb_epoch': 50,
+        'training': {
+            'batch_size': 100,
+            'nb_epoch': 2000,
             'validation_split': 0.1,
-            'optimizer': Adam(clipnorm=1e-2),
         },
 
-        'model_params': {
-            'n_embed_dims': 100,
-            'n_hidden': 200,
-
-            # convolution
-            'nb_filters': 1000,  # * 4
-            'conv_activation': 'tanh',
-
-            # recurrent
-            'n_lstm_dims': 141,  # * 2
-
-            'initial_embed_weights': np.load('models/word2vec_100_dim.h5'),
-            'similarity_dropout': 0.5,
-        },
-
-        'similarity_params': {
+        'similarity': {
             'mode': 'gesd',
             'gamma': 1,
             'c': 1,
@@ -252,28 +225,12 @@ if __name__ == '__main__':
         }
     }
 
-    evaluator = Evaluator(conf)
-
-    ##### Define model ######
-    model = AttentionModel(conf)
-    optimizer = conf.get('training_params', dict()).get('optimizer', 'rmsprop')
-    model.compile(optimizer=optimizer)
-
-    # save embedding layer
-    # evaluator.load_epoch(model, 7)
-    # embedding_layer = model.prediction_model.layers[2].layers[2]
-    # weights = embedding_layer.get_weights()[0]
-    # np.save(open('models/embedding_1000_dim.h5', 'wb'), weights)
+    from keras_models import ConvolutionModel
+    evaluator = Evaluator(conf, model=ConvolutionModel, optimizer=SGD(lr=0.001))
 
     # train the model
-    # evaluator.load_epoch(model, 6)
-    best_loss = evaluator.train(model)
+    best_loss = evaluator.train()
 
     # evaluate mrr for a particular epoch
-    evaluator.load_epoch(model, best_loss['epoch'])
-    # evaluator.load_epoch(model, 31)
-    evaluator.get_mrr(model, evaluate_all=True)
-    # for epoch in range(1, 100):
-    #     print('Epoch %d' % epoch)
-    #     evaluator.load_epoch(model, epoch)
-    #     evaluator.get_mrr(model, evaluate_all=True)
+    evaluator.load_epoch(best_loss['epoch'])
+    evaluator.get_score(verbose=False)
diff --git a/keras_models.py b/keras_models.py
index 27e3b32..fe7b0ed 100644
--- a/keras_models.py
+++ b/keras_models.py
@@ -3,14 +3,12 @@ from __future__ import print_function
 from abc import abstractmethod
 
 from keras.engine import Input
-from keras.layers import merge, Embedding, Dropout, Convolution1D, Lambda, LSTM, Dense, TimeDistributed, constraints
+from keras.layers import merge, Embedding, Dropout, Convolution1D, Lambda, LSTM, Dense
 from keras import backend as K
 from keras.models import Model
 
 import numpy as np
 
-from attention_lstm import AttentionLSTM
-
 
 class LanguageModel:
     def __init__(self, config):
@@ -19,8 +17,7 @@ class LanguageModel:
         self.answer_bad = Input(shape=(config['answer_len'],), dtype='int32', name='answer_bad_base')
 
         self.config = config
-        self.model_params = config.get('model_params', dict())
-        self.similarity_params = config.get('similarity_params', dict())
+        self.params = config.get('similarity', dict())
 
         # initialize a bunch of variables that will be set later
         self._models = None
@@ -41,14 +38,14 @@ class LanguageModel:
         return
 
     def get_similarity(self):
-        ''' Specify similarity in configuration under 'similarity_params' -> 'mode'
-        If a parameter is needed for the model, specify it in 'similarity_params'
+        ''' Specify similarity in configuration under 'similarity' -> 'mode'
+        If a parameter is needed for the model, specify it in 'similarity'
 
         Example configuration:
 
         config = {
             ... other parameters ...
-            'similarity_params': {
+            'similarity': {
                 'mode': 'gesd',
                 'gamma': 1,
                 'c': 1,
@@ -65,11 +62,11 @@ class LanguageModel:
         aesd: (euclidean + sigmoid) / 2
         '''
 
-        params = self.similarity_params
+        params = self.params
         similarity = params['mode']
 
         dot = lambda a, b: K.batch_dot(a, b, axes=1)
-        l2_norm = lambda a, b: K.sqrt(K.sum((a - b) ** 2, axis=1, keepdims=True))
+        l2_norm = lambda a, b: K.sqrt(K.sum(K.square(a - b), axis=1, keepdims=True))
 
         if similarity == 'cosine':
             return lambda x: dot(x[0], x[1]) / K.maximum(K.sqrt(dot(x[0], x[0]) * dot(x[1], x[1])), K.epsilon())
@@ -100,13 +97,11 @@ class LanguageModel:
 
         if self._qa_model is None:
             question_output, answer_output = self._models
-            dropout = Dropout(self.similarity_params.get('similarity_dropout', 0.2))
-            similarity = lambda x: K.expand_dims(self.get_similarity()(x), 1)
+            dropout = Dropout(self.params.get('similarity_dropout', 0.2))
+            similarity = self.get_similarity()
             qa_model = merge([dropout(question_output), dropout(answer_output)],
                              mode=similarity, output_shape=lambda _: (None, 1))
-                             # mode='cos', dot_axes=1)
-            self._qa_model = Model(input=[self.question, self.get_answer()], output=qa_model)
-            print(self._qa_model.output_shape)
+            self._qa_model = Model(input=[self.question, self.get_answer()], output=qa_model, name='qa_model')
 
         return self._qa_model
 
@@ -120,10 +115,10 @@ class LanguageModel:
                      mode=lambda x: K.relu(self.config['margin'] - x[0] + x[1]),
                      output_shape=lambda x: x[0])
 
-        self.prediction_model = Model(input=[self.question, self.answer_good], output=good_similarity)
+        self.prediction_model = Model(input=[self.question, self.answer_good], output=good_similarity, name='prediction_model')
         self.prediction_model.compile(loss=lambda y_true, y_pred: y_pred, optimizer=optimizer, **kwargs)
 
-        self.training_model = Model(input=[self.question, self.answer_good, self.answer_bad], output=loss)
+        self.training_model = Model(input=[self.question, self.answer_good, self.answer_bad], output=loss, name='training_model')
         self.training_model.compile(loss=lambda y_true, y_pred: y_pred, optimizer=optimizer, **kwargs)
 
     def fit(self, x, **kwargs):
@@ -150,11 +145,10 @@ class EmbeddingModel(LanguageModel):
         answer = self.get_answer()
 
         # add embedding layers
-        weights = self.model_params.get('initial_embed_weights', None)
+        weights = np.load(self.config['initial_embed_weights'])
         weights = weights if weights is None else [weights]
         embedding = Embedding(input_dim=self.config['n_words'],
-                              output_dim=self.model_params.get('n_embed_dims', 100),
-                              # W_constraint=constraints.nonneg(),
+                              output_dim=weights.shape[1],
                               weights=weights,
                               mask_zero=True)
         question_embedding = embedding(question)
@@ -169,8 +163,6 @@ class EmbeddingModel(LanguageModel):
 
 
 class ConvolutionModel(LanguageModel):
-    ### Validation loss at Epoch 65: 2.4e-6
-
     def build(self):
         assert self.config['question_len'] == self.config['answer_len']
 
@@ -178,39 +170,65 @@ class ConvolutionModel(LanguageModel):
         answer = self.get_answer()
 
         # add embedding layers
-        weights = self.model_params.get('initial_embed_weights', None)
+        weights = np.load(self.config['initial_embed_weights'])
         weights = weights if weights is None else [weights]
         embedding = Embedding(input_dim=self.config['n_words'],
-                              output_dim=self.model_params.get('n_embed_dims', 100),
+                              output_dim=weights.shape[1],
                               weights=weights)
         question_embedding = embedding(question)
         answer_embedding = embedding(answer)
 
-        # turn off layer updating
-        # embedding.params = []
-        # embedding.updates = []
+        # cnn
+        cnns = [Convolution1D(filter_length=filter_length,
+                              nb_filter=500,
+                              activation='tanh',
+                              border_mode='same') for filter_length in [2, 3, 5, 7]]
+        question_cnn = merge([cnn(question_embedding) for cnn in cnns], mode='concat')
+        answer_cnn = merge([cnn(answer_embedding) for cnn in cnns], mode='concat')
 
-        # dense
-        dense = TimeDistributed(Dense(self.model_params.get('n_hidden', 200),
-                                      # activity_regularizer=regularizers.activity_l1(1e-4),
-                                      # W_regularizer=regularizers.l1(1e-4),
-                                      activation='tanh'))
-        question_dense = dense(question_embedding)
-        answer_dense = dense(answer_embedding)
+        # maxpooling
+        maxpool = Lambda(lambda x: K.max(x, axis=1, keepdims=False), output_shape=lambda x: (x[0], x[2]))
+        enc = Dense(100, activation='tanh')
+        question_pool = enc(maxpool(question_cnn))
+        answer_pool = enc(maxpool(answer_cnn))
+
+        return question_pool, answer_pool
+
+
+class ConvolutionalLSTM(LanguageModel):
+    def build(self):
+        question = self.question
+        answer = self.get_answer()
+
+        # add embedding layers
+        weights = np.load(self.config['initial_embed_weights'])
+        weights = weights if weights is None else [weights]
+        embedding = Embedding(input_dim=self.config['n_words'],
+                              output_dim=weights.shape[1],
+                              weights=weights)
+        question_embedding = embedding(question)
+        answer_embedding = embedding(answer)
+
+        f_rnn = LSTM(141, return_sequences=True, consume_less='mem')
+        b_rnn = LSTM(141, return_sequences=True, consume_less='mem')
+
+        qf_rnn = f_rnn(question_embedding)
+        qb_rnn = b_rnn(question_embedding)
+        question_pool = merge([qf_rnn, qb_rnn], mode='concat', concat_axis=-1)
+
+        af_rnn = f_rnn(answer_embedding)
+        ab_rnn = b_rnn(answer_embedding)
+        answer_pool = merge([af_rnn, ab_rnn], mode='concat', concat_axis=-1)
 
         # cnn
         cnns = [Convolution1D(filter_length=filter_length,
-                              nb_filter=self.model_params.get('nb_filters', 1000),
-                              activation=self.model_params.get('conv_activation', 'relu'),
-                              # W_regularizer=regularizers.l1(1e-4),
-                              # activity_regularizer=regularizers.activity_l1(1e-4),
-                              border_mode='same') for filter_length in [2, 3, 5, 7]]
-        question_cnn = merge([cnn(question_dense) for cnn in cnns], mode='concat')
-        answer_cnn = merge([cnn(answer_dense) for cnn in cnns], mode='concat')
+                          nb_filter=500,
+                          activation='tanh',
+                          border_mode='same') for filter_length in [1, 2, 3, 5]]
+        question_cnn = merge([cnn(question_pool) for cnn in cnns], mode='concat')
+        answer_cnn = merge([cnn(answer_pool) for cnn in cnns], mode='concat')
 
-        # maxpooling
-        maxpool = Lambda(lambda x: K.max(x, axis=-1, keepdims=False), output_shape=lambda x: (x[0], x[2]))
-        avepool = Lambda(lambda x: K.mean(x, axis=-1, keepdims=False), output_shape=lambda x: (x[0], x[2]))
+        maxpool = Lambda(lambda x: K.max(x, axis=1, keepdims=False), output_shape=lambda x: (x[0], x[2]))
         question_pool = maxpool(question_cnn)
         answer_pool = maxpool(answer_cnn)
 
@@ -223,37 +241,31 @@ class AttentionModel(LanguageModel):
         answer = self.get_answer()
 
         # add embedding layers
-        weights = self.model_params.get('initial_embed_weights', None)
+        weights = np.load(self.config['initial_embed_weights'])
         weights = weights if weights is None else [weights]
         embedding = Embedding(input_dim=self.config['n_words'],
-                              output_dim=self.model_params.get('n_embed_dims', 100),
+                              output_dim=weights.shape[1],
                               weights=weights,
-                              mask_zero=True)
+                              # mask_zero=True)
+                              mask_zero=False)
         question_embedding = embedding(question)
         answer_embedding = embedding(answer)
 
-        # turn off layer updating
-        # embedding.params = []
-        # embedding.updates = []
-
         # question rnn part
-        f_rnn = LSTM(self.model_params.get('n_lstm_dims', 141), return_sequences=True, dropout_U=0.2,
-                     consume_less='mem')
-        b_rnn = LSTM(self.model_params.get('n_lstm_dims', 141), return_sequences=True, dropout_U=0.2,
-                     consume_less='mem', go_backwards=True)
+        f_rnn = LSTM(141, return_sequences=True, consume_less='mem')
+        b_rnn = LSTM(141, return_sequences=True, consume_less='mem', go_backwards=True)
         question_f_rnn = f_rnn(question_embedding)
         question_b_rnn = b_rnn(question_embedding)
 
         # maxpooling
         maxpool = Lambda(lambda x: K.max(x, axis=1, keepdims=False), output_shape=lambda x: (x[0], x[2]))
-        avepool = Lambda(lambda x: K.mean(x, axis=1, keepdims=False), output_shape=lambda x: (x[0], x[2]))
         question_pool = merge([maxpool(question_f_rnn), maxpool(question_b_rnn)], mode='concat', concat_axis=-1)
 
         # answer rnn part
-        f_rnn = AttentionLSTM(self.model_params.get('n_lstm_dims', 141), question_pool, return_sequences=True,
-                              consume_less='mem', single_attention_param=True)
-        b_rnn = AttentionLSTM(self.model_params.get('n_lstm_dims', 141), question_pool, return_sequences=True,
-                              consume_less='mem', go_backwards=True, single_attention_param=True)
+        from attention_lstm import AttentionLSTMWrapper
+        f_rnn = AttentionLSTMWrapper(f_rnn, question_pool, single_attention_param=True)
+        b_rnn = AttentionLSTMWrapper(b_rnn, question_pool, single_attention_param=True)
+
         answer_f_rnn = f_rnn(answer_embedding)
         answer_b_rnn = b_rnn(answer_embedding)
         answer_pool = merge([maxpool(answer_f_rnn), maxpool(answer_b_rnn)], mode='concat', concat_axis=-1)
diff --git a/results.notes b/results.notes
index bff088b..eebfa17 100644
--- a/results.notes
+++ b/results.notes
@@ -10,16 +10,6 @@ Embedding + Max Pooling:
    - 0.611 on test 2
    - 0.624 on dev
 
-Dense + CNN + Max Pooling:
- - Top 1 precision:
-   - 0.507 on test 1
-   - 0.458 on test 2
-   - 0.515 on dev
- - MRR:
-   - 0.635 on test 1
-   - 0.593 on test 2
-   - 0.642 on dev
-
 Attentional LSTM + Max Pooling:
  - Top 1 precision:
    - 0.480 on test 1
@@ -35,8 +25,17 @@ Unsupervised RNN language model + trained embeddings:
    - 0.546 on test 1
    - 0.527 on test 2
    - 0.552 on dev
- - Mrr:
+ - MRR:
    - 0.670 on test 1
    - 0.651 on test 2
    - 0.671 on dev
 
+Training ConvolutionalLSTM model for a long time (~4 days):
+ - Top-1 Precision:
+   - 0.564 on test 1
+   - 0.543 on test 2
+   - 0.573 on dev
+ - MRR:
+   - 0.681 on test 1
+   - 0.661 on test 2
+   - 0.686 on dev
diff --git a/seq2seq/__init__.py b/seq2seq/__init__.py
deleted file mode 100644
index e69de29..0000000
diff --git a/seq2seq/answer_to_question.py b/seq2seq/answer_to_question.py
deleted file mode 100644
index 13307c4..0000000
--- a/seq2seq/answer_to_question.py
+++ /dev/null
@@ -1,218 +0,0 @@
-'''
-Model for sequence to sequence learning. The model learns to generate a question given an answer,
-and generalizes to other questions and answers.
-'''
-
-from __future__ import print_function
-
-import os
-import random
-
-import numpy as np
-from keras.engine import Input
-from keras.layers import RepeatVector, TimeDistributed, Dense, Activation, merge, GRU, Embedding, regularizers, Lambda, \
-    constraints
-from keras.models import Model
-import keras.backend as K
-
-from keras_models import LanguageModel
-
-try:
-    import cPickle as pickle
-except:
-    import pickle
-
-data_path = os.environ['INSURANCE_QA']
-model_save = os.path.join(os.environ['MODEL_PATH'], 'model.h5')
-
-
-class InsuranceQA:
-    def __init__(self):
-        self.vocab = self.load('vocabulary')
-        self.table = InsuranceQA.VocabularyTable(self.vocab.values())
-
-    def load(self, name):
-        return pickle.load(open(os.path.join(data_path, name), 'rb'))
-
-    def save(self, obj, name):
-        pickle.dump(obj, open(os.path.join(data_path, name), 'wb'))
-
-    class VocabularyTable:
-        def __init__(self, words):
-            self.words = sorted(set(words))
-            self.words_indices = dict((c, i) for i, c in enumerate(self.words))
-            self.indices_words = dict((i, c) for i, c in enumerate(self.words))
-
-        def encode(self, sentence, maxlen, one_hot=False):
-            if one_hot:
-                indices = np.zeros((maxlen, len(self.words) + 1), dtype=np.int32)
-                for i, w in enumerate(sentence):
-                    if i == maxlen: break
-                    indices[i, self.words_indices[w]] = 1
-                return indices
-            else:
-                indices = np.zeros((maxlen,), dtype=np.int32)
-                for i, w in enumerate(sentence):
-                    if i == maxlen: break
-                    indices[i] = self.words_indices[w]
-                return indices
-
-        def decode(self, indices, calc_argmax=True):
-            if calc_argmax:
-                indices = np.argmax(indices, axis=-1)
-            return ' '.join(self.indices_words[x] for x in indices if x != 0)
-
-
-def get_model(question_maxlen, answer_maxlen, vocab_len, n_hidden, load_save=False):
-    answer = Input(shape=(answer_maxlen,), dtype='int32')
-    embedded = Embedding(input_dim=vocab_len, output_dim=n_hidden, mask_zero=True)(answer)
-    # answer = Input(shape=(answer_maxlen, vocab_len))
-    # embedded = Masking(mask_value=0.)(answer)
-
-    # encoder rnn
-    encode_rnn = GRU(n_hidden, return_sequences=True, dropout_U=0.2)(embedded)
-    encode_rnn = GRU(n_hidden, return_sequences=False, dropout_U=0.2)(encode_rnn)
-
-    encode_brnn = GRU(n_hidden, return_sequences=True, go_backwards=True, dropout_U=0.2)(embedded)
-    encode_brnn = GRU(n_hidden, return_sequences=False, go_backwards=True, dropout_U=0.2)(encode_brnn)
-
-    # repeat it maxlen times
-    repeat_encoding_rnn = RepeatVector(question_maxlen)(encode_rnn)
-    repeat_encoding_brnn = RepeatVector(question_maxlen)(encode_brnn)
-
-    # decoder rnn
-    decode_rnn = GRU(n_hidden, return_sequences=True, dropout_U=0.2, dropout_W=0.5)(repeat_encoding_rnn)
-    decode_rnn = GRU(n_hidden, return_sequences=True, dropout_U=0.2)(decode_rnn)
-
-    decode_brnn = GRU(n_hidden, return_sequences=True, go_backwards=True, dropout_U=0.2, dropout_W=0.5)(
-        repeat_encoding_brnn)
-    decode_brnn = GRU(n_hidden, return_sequences=True, go_backwards=True, dropout_U=0.2)(decode_brnn)
-
-    merged_output = merge([decode_rnn, decode_brnn], mode='concat', concat_axis=-1)
-
-    # output
-    dense = TimeDistributed(Dense(vocab_len, activity_regularizer=regularizers.activity_l1(1e-4)))(merged_output)
-    softmax = Activation('softmax')(dense)
-
-    # compile the prediction model
-    model = Model([answer], [softmax])
-
-    model.compile(loss='categorical_crossentropy', optimizer='rmsprop', metrics=['accuracy'])
-
-    if os.path.exists(model_save) and load_save:
-        model.load_weights(model_save)
-
-    return model
-
-
-class EmbeddingRNNModel(LanguageModel):
-    def build(self):
-        question = self.question
-        answer = self.get_answer()
-
-        rnn_model = get_model(question_maxlen=self.model_params.get('question_len', 20),
-                              answer_maxlen=self.model_params.get('question_len', 60),
-                              vocab_len=self.config['n_words'], n_hidden=256, load_save=True)
-        rnn_model.trainable = False
-
-        answer_inverted = rnn_model(answer)
-        argmax = Lambda(lambda x: K.argmax(x, axis=2), output_shape=lambda x: (x[0], x[1]))
-        argmax.trainable = False
-        answer_argmax = argmax(answer_inverted)
-
-        # add embedding layers
-        weights = self.model_params.get('initial_embed_weights', None)
-        weights = weights if weights is None else [weights]
-        embedding = Embedding(input_dim=self.config['n_words'],
-                              output_dim=self.model_params.get('n_embed_dims', 100),
-                              # W_regularizer=regularizers.activity_l1(1e-4),
-                              W_constraint=constraints.nonneg(),
-                              weights=weights,
-                              mask_zero=True)
-        question_embedding = embedding(question)
-        answer_embedding = embedding(answer_argmax)
-
-        # maxpooling
-        maxpool = Lambda(lambda x: K.max(x, axis=1, keepdims=False), output_shape=lambda x: (x[0], x[2]))
-        question_maxpool = maxpool(question_embedding)
-        answer_maxpool = maxpool(answer_embedding)
-
-        # activation
-        activation = Activation('linear')
-        question_output = activation(question_maxpool)
-        answer_output = activation(answer_maxpool)
-
-        return question_output, answer_output
-
-
-if __name__ == '__main__':
-    question_maxlen, answer_maxlen = 20, 60
-
-    qa = InsuranceQA()
-    batch_size = 50
-    n_test = 5
-    nb_epoch = 20
-    nb_iteration = 200
-
-    print('Generating data...')
-    answers = qa.load('answers')
-
-
-    def gen_questions(batch_size, test=False):
-        if test:
-            questions = qa.load('test1')
-        else:
-            questions = qa.load('train')
-        while True:
-            i = 0
-            question_idx = np.zeros(shape=(batch_size, question_maxlen, len(qa.vocab) + 1))
-            answer_idx = np.zeros(shape=(batch_size, answer_maxlen))
-            random.shuffle(questions)
-            for s in questions:
-                if test:
-                    ans = s['good']
-                else:
-                    ans = s['answers']
-                for a in ans:
-                    answer = qa.table.encode([qa.vocab[x] for x in answers[a]], answer_maxlen, one_hot=False)
-                    question = qa.table.encode([qa.vocab[x] for x in s['question']], question_maxlen, one_hot=True)
-                    # question = np.amax(question, axis=0, keepdims=False)
-                    answer_idx[i] = answer
-                    question_idx[i] = question
-                    i += 1
-                    if i == batch_size:
-                        yield ([answer_idx], [question_idx])
-                        i = 0
-
-
-    gen = gen_questions(batch_size)
-    test_gen = gen_questions(n_test, test=True)
-
-    print('Generating model...')
-    model = get_model(question_maxlen=question_maxlen, answer_maxlen=answer_maxlen, vocab_len=len(qa.vocab) + 1,
-                      n_hidden=256, load_save=True)
-
-    # print('Training model...')
-    for iteration in range(1, nb_iteration + 1):
-        print('\n' + '-' * 50 + '\nIteration %d' % iteration)
-        model.fit_generator(gen, samples_per_epoch=100 * batch_size, nb_epoch=nb_epoch)
-        model.save_weights(model_save, overwrite=True)
-
-        # test this iteration on some sample data
-        x, y = next(test_gen)
-        pred = model.predict(x, verbose=0)
-        y = y[0]
-        x = x[0]
-        for i in range(n_test):
-            print('Answer: {}'.format(qa.table.decode(x[i], calc_argmax=False)))
-            print('  Expected: {}'.format(qa.table.decode(y[i])))
-            print('  Predicted: {}'.format(qa.table.decode(pred[i])))
-
-    print('Saving data points...')
-    generated = dict()
-    for key, answer in answers.items():
-        print('\r%d / %d' % (key, len(answers)), end = '')
-        output = model.predict(qa.table.encode([qa.vocab[x] for x in answer], answer_maxlen, one_hot=False).reshape((1, answer_maxlen)))
-        argmax = np.argmax(output, axis=-1)[0]
-        generated[key] = answer
-    qa.save(generated, 'generated')
diff --git a/utils/__init__.py b/utils/__init__.py
deleted file mode 100644
index e69de29..0000000
diff --git a/utils/dictionary.py b/utils/dictionary.py
deleted file mode 100644
index 5d0ed88..0000000
--- a/utils/dictionary.py
+++ /dev/null
@@ -1,99 +0,0 @@
-from __future__ import print_function
-
-try:
-    import six.modes.cPickle as pickle
-except ImportError:
-    import pickle
-
-
-class Dictionary:
-    def __init__(self, min_len=1):
-        self._token_counts = dict()
-        self._id = 0
-        self._min_len = min_len
-
-        self.token2id = dict()
-        self.id2token = list()
-
-    def add(self, text):
-        if text is None: return
-
-        from gensim.utils import tokenize
-
-        if isinstance(text, str):
-            docs = [tokenize(text, to_lower=True)]
-        else:
-            docs = [tokenize(t, to_lower=True) for t in text]
-
-        for doc in docs:
-            for t in doc:
-                if t in self._token_counts:
-                    self._token_counts[t] += 1
-                else:
-                    self._token_counts[t] = 1
-                    self.id2token.append(t)
-                    self.token2id[t] = self._id
-                    self._id += 1
-
-    def __call__(self, item):
-        return self.token2id.get(item, self._id)
-
-    def __getitem__(self, item):
-        return self.id2token[item] if item < self._id else 'X'
-
-    def __len__(self):
-        return self._id + 1
-
-    def convert(self, text):
-        from gensim.utils import tokenize
-        from numpy import asarray
-
-        if isinstance(text, str):
-            docs = [tokenize(text, to_lower=True, deacc=True)]
-        else:
-            docs = [tokenize(t, to_lower=True, deacc=True) for t in text]
-
-        return [asarray([self(t) for t in doc], dtype='int32') for doc in docs]
-
-    def revert(self, tokens):
-        texts = list()
-
-        for token in tokens:
-            texts.append(' '.join([self[t] for t in token]))
-
-        return texts
-
-    def top(self, n):
-        import operator
-
-        sorted_tokens = sorted(self._token_counts.items(), reverse=True, key=operator.itemgetter(1))[:n]
-        self._token_counts = dict((k, v) for k, v in sorted_tokens)
-        self.id2token = [k for k in self._token_counts.keys()]
-        self.token2id = dict((v, k) for k, v in enumerate(self.id2token))
-        self._id = len(self.id2token)
-
-    def strip(self, n):
-        self._token_counts = dict((k, v) for k, v in self._token_counts.items() if v > n)
-        self.id2token = [k for k in self._token_counts.keys()]
-        self.token2id = dict((v, k) for k, v in enumerate(self.id2token))
-        self._id = len(self.id2token)
-
-    def save(self, file_name):
-        pickle.dump(self, open(file_name, 'wb+'))
-
-    def __repr__(self):
-        return '<Dictionary (%d tokens)>' % self._id
-
-    @staticmethod
-    def load(file_name):
-        return pickle.load(open(file_name, 'rb'))
-
-if __name__ == '__main__':
-    d = Dictionary()
-    d.add('the apples and oranges are very fresh today')
-    print(d)
-
-    c = d.convert('today, i want the fresh apples and oranges')
-    print(c)
-    r = d.revert(c)
-    print(r)