EXAMPLE : Interpreter grounded Neural Program Synthesis [WIP]

[reshinthadithyan]

This PR adds an example of using trlx to ground LMs into the interpreter with a toy list manipulation DSL to make them generate more coherent programs to satisfy IO conditions.

Questions:

1. I want to do extensive ablation studies on the impact of grounding Models with Interpreter. So would need a folder to add plots and ReadME.md. Does the current directory structure look good ?

Interpreter Grounded Program Synthesis

Program synthesis is the task of automatically generating programs that solve a given task by satisfying an IO condition. In Neural Program Synthesis the synthesizer is a neural network which is a Language Model that takes in an input/output pair and tries to generate the program in the defined toy DSL's Grammar.

Toy List Manipulation DSL Grammar

The DSL has the following grammar:

list_expr := list[int]
integer := -5 | -4 | -3 | -2 | -1 | 0 | 1 | 2 | 3 | 4 | 5
statement :=
          | take(list_expr,integer)
          | drop(list_expr,integer)
          | reverse(list_expr)
          | sort_asc(list_expr)
          | sort_des(list_expr)
          | add_n(list_expr,integer)
          | sub_n(list_expr,integer)
          | mul_n(list_expr,integer)
          | expand_copy(list_expr)

To generate training/testing data run, python3 -m lang. The dataset would be saved in ./dataset/train.json and ./dataset/test.json. To use the processed dataset refer to this google drive link.
Each datapoint in the dataset would look like,

    {"input": "Input: [4, -2, 0, 0, 5, 5] Output: [25, 25, 20, 0, 0, -10] Function:",
    "output": "sort_des(reverse(mul_n(sort_asc(sort_asc([4, -2, 0, 0, 5, 5])),5)))"}

Training with TRLX

Run python3 -m train_trlx.py to run the training with grounded interpreter. The reward_fn, would return -1 if a sample generated is of invalid syntax. it would return 0.5 if the generated syntax is valid but doesn't satisfy IO condition.

[cat-state]

thanks for this example @reshinthadithyan ! it mostly looks good aside from a few things. Could you also format the python files with black?

[cat-state]

Could you comment this file more? Maybe to separate the sections and an explanation at the top for the structure of the language/interpreter

[LouisCastricato]

Strong agree. The documentation here is not clear.

[reshinthadithyan]

Adding more documentation, thanks for the review.

[reshinthadithyan]

Added more documentation on the sampler.

[cat-state]

In addition to the DSL grammar, add some snippets

[reshinthadithyan]

Added example snippets to showcase the atomic functions.

[LouisCastricato]

This looks fantastic! I think we should upload a trained model to hugging face hub that goes along with this tutorial, and allow people to validate their model against our uploaded one. What do you think Reshinth?

[reshinthadithyan]

@LouisCastricato Yes, the idea is to give a fully reproducible script for the anyone to train it by themselves and validate.

[LouisCastricato]

LGTM, ready to merge with your approval @cat-state

[cat-state]

LGTM! thanks @reshinthadithyan !

[LouisCastricato]

Precommit is failing. Will merge after resolution.