Rewrite to use TF 1.1.0, support ML engine, updated results

Author: jimfleming

.dockerignore

@@ -0,0 +1,2 @@
+data/
+logs/

.editorconfig

@@ -1,4 +1,4 @@
-datasets/
+data/
 logs/
 .env/
 *.pyc

.gitignore

@@ -1,9 +1,9 @@
 # Recurrent Entity Networks
 
 This repository contains an independent TensorFlow implementation of recurrent entity networks from [Tracking the World State with
-Recurrent Entity Networks](https://openreview.net/forum?id=rJTKKKqeg). This paper introduces the first method to solve all of the bAbI tasks using 10k training examples. The author's original Torch implementation is available [here](https://github.com/facebook/MemNN/tree/master/EntNet-babi).
+Recurrent Entity Networks](https://arxiv.org/abs/1612.03969). This paper introduces the first method to solve all of the bAbI tasks using 10k training examples. The author's original Torch implementation is now available [here](https://github.com/facebook/MemNN/tree/master/EntNet-babi).
 
-<img src="images/diagram.png" alt="Diagram of recurrent entity network" width="886" height="658">
+<img src="assets/diagram.png" alt="Diagram of recurrent entity network architecture" width="886" height="658">
 
 ## Results
 
@@ -16,37 +16,40 @@ Task | EntNet (paper) | EntNet (repo)
 3: 3 supporting facts | 4.1 | ?
 4: 2 argument relations | 0 | 0
 5: 3 argument relations | 0.3 | ?
-6: yes/no questions | 0.2 | 0.1
-7: counting | 0 | ?
-8: lists/sets | 0.5 | ?
-9: simple negation | 0.1 | 0.7
-10: indefinite knowledge | 0.6 | 0.1
+6: yes/no questions | 0.2 | 0
+7: counting | 0 | 0
+8: lists/sets | 0.5 | 0
+9: simple negation | 0.1 | 0
+10: indefinite knowledge | 0.6 | 0
 11: basic coreference | 0.3 | 0
 12: conjunction | 0 | 0
 13: compound coreference | 1.3 | 0
-14: time reasoning | 0 | 4.5
+14: time reasoning | 0 | 0
 15: basic deduction | 0 | 0
-16: basic induction | 0.2 | 54.0 ([#5](../../issues/5))
+16: basic induction | 0.2 | 0
 17: positional reasoning | 0.5 | 1.7
 18: size reasoning | 0.3 | 1.5
-19: path finding | 2.3 | 41.9 ([#5](../../issues/5))
-20: agents motivation | 0 | 0.2
+19: path finding | 2.3 | 0
+20: agents motivation | 0 | 0
 **Failed Tasks** | 0 | ?
 **Mean Error** | 0.5 | ?
 
+NOTE: Some of these tasks (16 and 19, in particular) required a change in learning rate schedule to reliably converge.
+
 ## Setup
 
-1. Download the datasets by running [download_datasets.sh](download_datasets.sh) or from [The bAbI Project](https://research.facebook.com/research/babi/).
-2. Run [prep_datasets.py](prep_datasets.py) which will convert the datasets into [TFRecords](https://www.tensorflow.org/versions/r0.11/how_tos/reading_data/index.html#standard_tensorflow_format).
+1. Download the datasets by running [download_babi.sh](download_babi.sh) or from [The bAbI Project](https://research.facebook.com/research/babi/).
+2. Run [prep_data.py](entity_networks/prep_data.py) which will convert the datasets into [TFRecords](https://www.tensorflow.org/programmers_guide/reading_data#standard_tensorflow_format).
 3. Run `python -m entity_networks.main` to begin training on QA1.
-4. Run `./run_all.sh` to train on all tasks.
 
-## Dependencies
+## Major Dependencies
+
+- TensorFlow v1.1.0
 
-- TensorFlow v0.11
+(For additional dependencies see [requirements.txt](requirements.txt))
 
 ## Thanks!
 
 - Thanks to Mikael Henaff for providing details about their paper over Thanksgiving break. :)
 - Thanks to Andy Zhang ([@zhangandyx](https://twitter.com/zhangandyx)) for helping me troubleshoot numerical instabilities.
-- Thanks to Mike Young for providing results on some of the longer tasks.
+- Thanks to Mike Young (@mikalyoung) for providing results on some of the longer tasks.

README.md

@@ -0,0 +1,11 @@
+#!/bin/bash
+
+if [ ! -d ./datasets ]; then
+  mkdir -p ./datasets
+fi
+
+BABI_TASKS=datasets/babi_tasks_data_1_20_v1.2.tar.gz
+
+if [ ! -f $BABI_TASKS ]; then
+  wget http://www.thespermwhale.com/jaseweston/babi/tasks_1-20_v1-2.tar.gz -O $BABI_TASKS
+fi

images/diagram.png assets/diagram.png

@@ -0,0 +1,91 @@
+from __future__ import absolute_import
+from __future__ import print_function
+from __future__ import division
+
+import os
+import json
+import random
+import argparse
+import tensorflow as tf
+
+from tqdm import tqdm
+
+from entity_networks.inputs import generate_input_fn
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        '--data-dir',
+        help='Directory containing data',
+        default='data/babi/records/')
+    args = parser.parse_args()
+
+    tasks_dir = 'tasks/'
+
+    if not os.path.exists(tasks_dir):
+        os.makedirs(tasks_dir)
+
+    task_names = [
+        'qa1_single-supporting-fact',
+        'qa2_two-supporting-facts',
+        'qa3_three-supporting-facts',
+        'qa4_two-arg-relations',
+        'qa5_three-arg-relations',
+        'qa6_yes-no-questions',
+        'qa7_counting',
+        'qa8_lists-sets',
+        'qa9_simple-negation',
+        'qa10_indefinite-knowledge',
+        'qa11_basic-coreference',
+        'qa12_conjunction',
+        'qa13_compound-coreference',
+        'qa14_time-reasoning',
+        'qa15_basic-deduction',
+        'qa16_basic-induction',
+        'qa17_positional-reasoning',
+        'qa18_size-reasoning',
+        'qa19_path-finding',
+        'qa20_agents-motivations',
+    ]
+
+    for task_name in tqdm(task_name.iteritems()):
+        metadata_path = os.path.join(args.data_dir, '{}_10k.json'.format(task_name))
+        with open(metadata_path) as metadata_file:
+            metadata = json.load(metadata_file)
+
+        filename = os.path.join(data_dir, '{}_10k_{}.tfrecords'.format(dataset_id, 'test'))
+        input_fn = generate_input_fn(
+            filename=eval_filename,
+            metadata=metadata,
+            batch_size=BATCH_SIZE,
+            num_epochs=1,
+            shuffle=False)
+
+        with tf.Graph().as_default():
+            features, answer = input_fn()
+
+            story = features['story']
+            query = features['query']
+
+            instances = []
+
+            with tf.train.SingularMonitoredSession() as sess:
+                while not sess.should_stop():
+                    story_, query_, answer_ = sess.run([story, query, answer])
+
+                    instance = {
+                        'story': story_[0].tolist(),
+                        'query': query_[0].tolist(),
+                        'answer': answer_[0].tolist(),
+                    }
+
+                    instances.append(instance)
+
+            metadata['instances'] = random.sample(instances, k=10)
+
+            output_path = os.path.join(tasks_dir, '{}.json'.format(task_name))
+            with open(output_path, 'w') as f:
+                f.write(json.dumps(metadata))
+
+if __name__ == '__main__':
+    main()