This repository is the official implementation of the paper https://arxiv.org/abs/2101.04117, containing training code, evaluation scripts, and figure generation.

For an easy-to-use API implementing a trained version of this model, check out spock!

To train the model, download the data into the /data folder from this Globus link.

Run the script train.sh, which will train the model 30 times from different seeds.

To generate the figures, edit the figures/generate.sh script to have true instead of false for any figure you'd like to generate. Then, execute that script in the folder.

One can install a working conda environment from the environment.yml file:

conda env create -f environment.yml

Note: You may need to delete the line:

  - conda-forge::cudatoolkit=10.1

from the environment.yml file depending on your system setup.

Note 2: On Apple Silicon, in addition to the above change, you will also need to run the env create command with CONDA_SUBDIR=osx-64 as an environment variable. This is because some of the specific dependencies are not available for Apple Silicon yet.