PyTorch implementation of RobustFill

Original Paper: https://arxiv.org/pdf/1703.07469.pdf

Model Checkpoint: https://huggingface.co/eddyyeo/robustfill

The RobustFill network by Devlin et al. is trained for the following task -- based on a few example input-output string pairs, generate a program in a domain-specific language that transforms the given inputs into the given outputs. This program can then be used to transform unseen inputs. For example:

Given these pairs:

Input	Output
Jacob Devlin	Devlin, J.
Eddy Yeo	Yeo, E.
Andrej Karpathy	Karpathy, A.
Anatoly Yakovenko	Yakovenko, A.

The RobustFill network will generate a program that can be used to transform an unbounded number of unseen inputs:

Unseen input	Transformed Output
Elon Musk	Musk, E.
Joe Rogan	Rogan, J.
Balaji Srinivasan	Srinivasan, B.

The program generated by our trained network for the example above is as follows:

Concat(
    Compose(
        Trim(),
        GetFrom(<Type.LOWER: 6>)
    ),
    ConstStr(','),
    ConstStr(' '),
    GetUpto(<Type.CHAR: 8>),
    ConstStr('.')
)

See the demo notebook to reproduce the result with the model checkpoint:

• Demo Notebook
• Model Checkpoint

The network was trained on Google Cloud with 4 x NVIDIA Tesla P4 using PyTorch's Distributed Data Parallel.

Instructions

Set up environment:

python3 -m venv env
source env/bin/activate
pip install -r requirements.txt

Train neural net. The script will automatically use GPU(s) if they are available.

python train.py --mode full

For testing purposes, run smaller network (on CPU) with a smaller problem size just to see that the loss goes to 0.

python train.py --mode easy

Run profiler:

python train.py --mode profile

Run unit tests:

python -m unittest

Lint:

flake8