Credits:
framework, models/layers/geometric, models/layers/ndr_geometric_stack.py adapted from: https://github.com/RobertCsordas/ndr
models/layers/CRvNN.py adapted from: https://github.com/JRC1995/Continuous-RvNN
models/layers/OrderedMemory.py adapted from: https://github.com/yikangshen/Ordered-Memory
models/layers/GumbelTreeLSTM.py, models/layers/GumbelTreeCell.py adapted from: https://github.com/jihunchoi/unsupervised-treelstm
optimizers/ranger.py adapted from: https://github.com/anoidgit/transformer/blob/master/optm/ranger.py
models/topk_module.py taken from the supplementary code in: https://proceedings.neurips.cc/paper/2020/hash/ec24a54d62ce57ba93a531b460fa8d18-Abstract.html
We also used https://github.com/nshepperd/gumbel-rao-pytorch as a reference for some of our implementations. 

Requirements: Same as https://github.com/JRC1995/Continuous-RvNN

Data Setup
Put the Logical Inference data files (train0,train1,train2,.....test12) (downloaded from https://github.com/yikangshen/Ordered-Memory/tree/master/data/propositionallogic) in data/proplogic/
Download the ListOps data (along with extrapolation data) from the urls here: https://github.com/facebookresearch/latent-treelstm/blob/master/data/listops/external/urls.txt and put the tsv files in data/listops/
Download LRA (https://github.com/google-research/long-range-arena) dataset, put basic_test.tsv LRA listops test set on data/listops
Run all the make*.py files in data/listops/ to create relevant splits (exact splits used in the paper will be released later) 
SST2 and SST5 data is downloaded and used through torchtext (so you don't have to do anything for SST5/2 data at this point)
Download IMDB original split from here: http://ai.stanford.edu/~amaas/data/sentiment/aclImdb_v1.tar.gz, extract and put the acllmdb folder in data/IMDB/ 
Download IMDB contrast set from here: https://github.com/allenai/contrast-sets/tree/main/IMDb/data, put the dev_contrast.tsv and test_contrast.tsv in data/IMDB/
Download IMDB contrast set from here: https://github.com/allenai/contrast-sets/tree/main/IMDb/data, put the test_counterfactual.tsv in data/IMDB/
Download IMDB counterfactual test set from here: https://github.com/acmi-lab/counterfactually-augmented-data/blob/master/sentiment/new/test.tsv; rename it to "test_counterfactual.tsv". Put it in data/IMDB
Download MNLI datasets (https://cims.nyu.edu/~sbowman/multinli/) and stress tests (https://abhilasharavichander.github.io/NLI_StressTest/); put them in data/MNLI/
Put glove.840B.300d.txt (download glove.840B.300d.zip from here:https://nlp.stanford.edu/projects/glove/) in embeddings/glove/

Preprocess
Run the preprocess files in preprocess to preprocess relevant files (listopsc corresponds to original listops, listopsd corresponds to listops-DG for depth generalization tests, listops_ndr50 stands for listops-DG1, listops_ndr100 stands for listops-DG2

Train:
python trian.py --model=[insert model name] -- dataset=[insert dataset name] --times=[insert total runs] --device=[insert device name] --model_type=[classifier/sentence_pair]
Check argparser.py for exact options. 
(We use "Cell" to mean "GRC". BeamTreeGRC=StochasticBeamTreeCell, DiffBeamTreeCell = BT-GRC+Softpath, GumbelDiffBeamTreeCell = BT-GRC+Gumbelpath, Beam1GumbelTreeCell = Gumbel-BT-GRC, DiffSortBeamTreeCell = BT-GRC+SoftTopK)