diff --git a/.gitignore b/.gitignore
index 9dcf2f1..c4a9da5 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,3 +1,7 @@
+# data / models (also potentially very large)
+data/
+models/
+
 # pyc files aren't necessary
 *.pyc
 
diff --git a/insuranceqa.py b/insuranceqa.py
index 7c56274..08b8bf6 100644
--- a/insuranceqa.py
+++ b/insuranceqa.py
@@ -11,8 +11,9 @@ from scipy.stats import rankdata
 random.seed(42)
 
 data_path = '/media/moloch/HHD/MachineLearning/data/insuranceQA'
+models_path = 'models/'
 
-emb_d = pickle.load(open('word2vec.dict', 'rb'))
+emb_d = pickle.load(open(os.path.join(models_path, 'word2vec.dict'), 'rb'))
 rev_d = dict([(v, k) for k, v in emb_d.items()])
 
 with open(os.path.join(data_path, 'vocabulary'), 'r') as f:
@@ -55,6 +56,7 @@ def get_eval(f_name):
 
     q_data = list()
     a_data = list()
+    n_good = list()
 
     for qa_pair in lines.split('\n'):
         if len(qa_pair) == 0: continue
@@ -62,13 +64,14 @@ def get_eval(f_name):
         a, q, g = qa_pair.split('\t')
 
         good_answers = [int(i) for i in a.strip().split(' ')]
-        all_answers = set([int(i) for i in g.strip().split(' ') if i not in good_answers])
+        all_answers = [int(i) for i in g.strip().split(' ')]
 
         question = convert_from_idxs(q)
         q_data.append(pad_sequences([question], maxlen=maxlen, padding='post', truncating='post', value=22295))
-        a_data.append(pad_sequences([answers[i] for i in [good_answers[0]] + list(all_answers)], maxlen=maxlen, padding='post', truncating='post', value=22295))
+        a_data.append(pad_sequences([answers[i] for i in all_answers], maxlen=maxlen, padding='post', truncating='post', value=22295))
+        n_good.append(len(good_answers))
 
-    return q_data, a_data
+    return q_data, a_data, n_good
 
 
 def get_data(f_name):
@@ -126,11 +129,12 @@ def get_accurate_percentage(model, questions, good_answers, bad_answers, n_eval=
     return correct
 
 
-def get_mrr(model, questions, all_answers, n_eval=512):
+def get_mrr(model, questions, all_answers, n_good, n_eval=-1):
 
-    if n_eval != 'all':
+    if n_eval != -1:
         questions = questions[-n_eval:]
         all_answers = all_answers[-n_eval:]
+        n_good = n_good[-n_eval:]
 
     c = 0
 
@@ -143,15 +147,7 @@ def get_mrr(model, questions, all_answers, n_eval=512):
         sims = model.predict([qs, ans]).flatten()
         r = rankdata(sims)
 
-        print(i)
-        print(revert(answers[np.argmax(r)]))
-        print(revert(question[0]))
-
-        print(sims)
-        print(r)
-
-        x = 1 / float(max(r) - r[0] + 1)
-        print(x)
+        x = 1 / float(max(r) - max(r[:n_good[i]]) + 1)
 
         c += x
 
@@ -173,20 +169,22 @@ data_sets = [
     'question.test1.label.token_idx.pool',
     'question.test2.label.token_idx.pool',
 ]
-# d_set = data_sets[1]
-# q_data, ag_data, ab_data, targets = get_data(d_set)
+q_data, ag_data, ab_data, targets = get_data(data_sets[0])
+qv_data, avg_data, avb_data, v_targets = get_data(data_sets[1])
 
 # found through experimentation that ~24 epochs generalized the best
-# print('Fitting model')
-# train_model.fit([q_data, ag_data, ab_data], targets, nb_epoch=24, batch_size=128, validation_split=0.2)
-# train_model.save_weights('iqa_model_for_training.h5', overwrite=True)
-# test_model.save_weights('iqa_model_for_prediction.h5', overwrite=True)
+print('Fitting model')
+for i in range(100):
+    np.random.shuffle(ab_data)
+    train_model.fit([q_data, ag_data, ab_data], targets, nb_epoch=1, batch_size=128, validation_data=[[qv_data, avg_data, avb_data], v_targets], shuffle=True)
+
+train_model.save_weights(os.path.join(models_path, 'iqa_model_for_training.h5'), overwrite=True)
+test_model.save_weights(os.path.join(models_path, 'iqa_model_for_prediction.h5'), overwrite=True)
 
 # the model actually did really well, predicted correct vs. incorrect answer 85% of the time on the validation set
-# test_model.load_weights('iqa_model_for_prediction.h5')
+# test_model.load_weights(os.path.join(models_path, 'iqa_model_for_prediction.h5'))
 # print('Percent correct: {}'.format(get_accurate_percentage(test_model, q_data, ag_data, ab_data, n_eval='all')))
 
-d_set = data_sets[1]
-q_data, a_data = get_eval(d_set)
-test_model.load_weights('iqa_model_for_prediction.h5')
-print('MRR: {}'.format(get_mrr(test_model, q_data, a_data)))
\ No newline at end of file
+q_data, a_data, n_good = get_eval(data_sets[1])
+# test_model.load_weights(os.path.join(models_path, 'iqa_model_for_prediction.h5'))
+print('MRR: {}'.format(get_mrr(test_model, q_data, a_data, n_good)))
diff --git a/keras_attention_model.py b/keras_attention_model.py
index b11609a..2b9572d 100644
--- a/keras_attention_model.py
+++ b/keras_attention_model.py
@@ -3,17 +3,21 @@ from __future__ import print_function
 ##############
 # Make model #
 ##############
+import os
+
 from keras.engine import Merge
 from keras.layers import Lambda, MaxPooling1D, Dense, Flatten, Dropout, Masking, Embedding, TimeDistributed
 from keras.optimizers import SGD
 
-from language_model.word_embeddings import Word2VecEmbedding
+from word_embeddings import Word2VecEmbedding
+
+models_path = 'models/'
 
 
 def make_model(maxlen, n_words, n_lstm_dims=141, n_embed_dims=128):
     from keras.optimizers import RMSprop
 
-    from language_model.attention_lstm import AttentionLSTM
+    from attention_lstm import AttentionLSTM
 
     from keras.layers import Input, LSTM, merge
     from keras.models import Model
@@ -25,8 +29,8 @@ def make_model(maxlen, n_words, n_lstm_dims=141, n_embed_dims=128):
     answer_bad = Input(shape=(maxlen,), dtype='int32')
 
     # language model
-    embedding = Embedding(n_words, n_embed_dims)
-    # embedding = Word2VecEmbedding('word2vec.model')
+    # embedding = Embedding(n_words, n_embed_dims)
+    embedding = Word2VecEmbedding(os.path.join(models_path, 'word2vec.model'))
 
     # forward and backward lstms
     f_lstm = LSTM(n_lstm_dims, return_sequences=True)
@@ -99,7 +103,7 @@ if __name__ == '__main__':
     # get the data set
     maxlen = 40 # words
 
-    from language_model.get_data import get_data_set, create_dictionary_from_qas
+    from utils.get_data import get_data_set, create_dictionary_from_qas
 
     dic = create_dictionary_from_qas()
     targets, questions, good_answers, bad_answers, n_dims = get_data_set(maxlen)
@@ -109,4 +113,4 @@ if __name__ == '__main__':
 
     print('Fitting model')
     train_model.fit([questions, good_answers, bad_answers], targets, nb_epoch=5, batch_size=32, validation_split=0.2)
-    train_model.save_weights('attention_lm_weights.h5', overwrite=True)
+    train_model.save_weights(os.path.join(models_path, 'attention_lm_weights.h5'), overwrite=True)
diff --git a/results.notes b/results.notes
new file mode 100644
index 0000000..e5eedf1
--- /dev/null
+++ b/results.notes
@@ -0,0 +1,4 @@
+Single-layer bi-LSTM with max pooling, 40 words per sentence, loss margin of 0.2
+ - MRR ~0.17
+ - Seemed to converge after about 20 epochs, with randomization between epochs
+
diff --git a/utils/dictionary.py b/utils/dictionary.py
index eca5c7a..b8ec234 100644
--- a/utils/dictionary.py
+++ b/utils/dictionary.py
@@ -5,8 +5,6 @@ try:
 except ImportError:
     import pickle
 
-import sys
-
 from numpy import asarray
 import numpy as np
 from gensim.utils import tokenize
diff --git a/utils/get_data.py b/utils/get_data.py
index 06c0919..7f0a328 100644
--- a/utils/get_data.py
+++ b/utils/get_data.py
@@ -6,18 +6,19 @@ import os
 import numpy as np
 from keras.preprocessing.sequence import pad_sequences
 
-from language_model.dictionary import Dictionary
+from utils.dictionary import Dictionary
 
 rng = np.random.RandomState(42)
 
+models_path = 'models/'
 file_name = 'liveqa-2015-rels.txt'
-dict_path = 'dict.pkl'
+dict_path = os.path.join(models_path, 'dict.pkl')
+
 
 #################
 # Download data #
 #################
 
-
 def download_data():
     from tqdm import tqdm
     import requests
@@ -25,20 +26,20 @@ def download_data():
     url = 'https://raw.githubusercontent.com/codekansas/ml/master/theano_stuff/ir/LiveQA2015-qrels-ver2.txt'
     response = requests.get(url, stream=True)
 
-    with open(file_name, 'wb') as handle:
+    with open(os.path.join(models_path, file_name), 'wb') as handle:
         for data in tqdm(response.iter_content()):
             handle.write(data)
 
+
 #################
 # Load QA pairs #
 #################
 
-
 def load_qa_pairs():
-    if not os.path.exists(file_name):
+    if not os.path.exists(os.path.join(models_path, file_name)):
         download_data()
 
-    with open(file_name, 'r') as f:
+    with open(os.path.join(models_path, file_name), 'r') as f:
         lines = re.split('\n|\r', f.read()) # for cross-system compatibility (not sure about windows)
 
     questions = dict()
@@ -65,6 +66,7 @@ def load_qa_pairs():
 
     return questions, answers
 
+
 #####################
 # Create dictionary #
 #####################
@@ -89,11 +91,11 @@ def create_dictionary_from_qas(questions=None, answers=None):
 
     return dic
 
+
 #################
 # Training set #
 ################
 
-
 def get_data_set(maxlen, questions=None, answers=None, dic=None):
 
     if questions is None or answers is None:
@@ -118,8 +120,8 @@ def get_data_set(maxlen, questions=None, answers=None, dic=None):
         targets += [0] * m
 
     targets = np.asarray(targets)
-    qs = pad_sequences(qs, maxlen=maxlen, padding='post', truncating='post')
-    gans = pad_sequences(gans, maxlen=maxlen, padding='post', truncating='post')
-    bans = pad_sequences(bans, maxlen=maxlen, padding='post', truncating='post')
+    qs = pad_sequences(qs, maxlen=maxlen, padding='post', truncating='post', dtype='int32')
+    gans = pad_sequences(gans, maxlen=maxlen, padding='post', truncating='post', dtype='int32')
+    bans = pad_sequences(bans, maxlen=maxlen, padding='post', truncating='post', dtype='int32')
 
     return targets, qs, gans, bans, len(dic)
diff --git a/word_embeddings.py b/word_embeddings.py
index cc67473..1e8348f 100644
--- a/word_embeddings.py
+++ b/word_embeddings.py
@@ -1,16 +1,20 @@
 import os
+
+import operator
 from gensim.models import Word2Vec
 from keras.engine import Layer
 import pickle
 
 import keras.backend as K
 
+models_path = 'models/'
+
 
 class Word2VecEmbedding(Layer):
     def __init__(self, model_path, **kwargs):
-        model = Word2Vec.load(model_path)
-        self.W = K.variable(model.syn0)
-        self.model_dims = model.syn0.shape
+        self.model = Word2Vec.load(model_path)
+        self.W = K.variable(self.model.syn0)
+        self.model_dims = self.model.syn0.shape
         super(Word2VecEmbedding, self).__init__(**kwargs)
 
     def build(self, input_shape):
@@ -21,6 +25,7 @@ class Word2VecEmbedding(Layer):
         return (input_shape[0], input_shape[1], self.model_dims[1])
 
     def call(self, x, mask=None):
+        x = K.maximum(K.minimum(x, self.model_dims[1] - 1), 0)
         return K.gather(self.W, x)
 
 
@@ -74,16 +79,20 @@ def train_model():
     sentences = questions + answers
 
     model = Word2Vec(sentences, size=100, min_count=1)
-    model.save('word2vec.model')
+    model.save(os.path.join(models_path, 'word2vec.model'))
 
 if __name__ == '__main__':
+    print('Training word2vec model..')
+
     train_model()
-    model = Word2Vec.load('word2vec.model')
+    model = Word2Vec.load(os.path.join(models_path, 'word2vec.model'))
+
+    print('Done! Saving...')
 
     d = dict([(k, v.index) for k, v in model.vocab.items()])
-    pickle.dump(d, open('word2vec.dict', 'wb'))
+    pickle.dump(d, open(os.path.join(models_path, 'word2vec.dict'), 'wb'))
+
+    d = pickle.load(open(os.path.join(models_path, 'word2vec.dict'), 'rb'))
+    print(sorted(d.items(), key=operator.itemgetter(1)))
 
-    # d = pickle.load(open('word2vec.dict', 'rb'))
-    # print(sorted(list(d.items()), key=lambda a, b: b))
-    #
     # Use the dictionary to convert sentences to vectors for dataset