Code generation with reinforcement learning
A code generation system uses a non-terminal expansion model and a non-terminal selector model to generate a code sketch to complete a partially-formed source code snippet. The non-terminal expansion model is a neural transformer model trained on a supervised dataset through reinforcement learning to learn to predict the production rule to expand for a given non-terminal symbol. The non-terminal selector model is trained through reinforcement learning to predict the non-terminal symbol to expand given a partial-code state. The models are used in a two-step beam search to generate the top candidate code sketches, where a candidate code sketch may contain a hole that represents an unexpanded non-terminal symbol.
1. A system for grammar-based code generation comprising:
a processor and a memory;
wherein the memory includes executable instructions that when executed by the processor perform actions that:
obtain a partially-formed source code snippet of a source code program associated with a programming language, wherein the programming language is associated with a grammar comprising a plurality of production rules, a set of non-terminal symbols, and a set of terminal symbols, wherein a production rule replaces a non-terminal symbol with one or more terminal symbols or one or more non-terminal symbols;
transform the partially-formed source code snippet into a partial code state comprising at least one non-terminal symbol; and
generate a code sketch to complete the partially-formed source code snippet, wherein the code sketch comprises an output sequence of terminal symbols and at least one hole,
wherein the generation of the code sketch uses a non-terminal selector model to predict a select non-terminal symbol in the partial code state to expand and a non-terminal expansion model to predict a select production rule to expand the select non-terminal symbol, and
wherein the hole represents a position in the output sequence where the non-terminal selector model is uncertain about which non-terminal symbol to expand.
2. The system of claim 1 , wherein the memory includes executable instructions that when executed by the processor perform actions that:
output the code sketch in the source code program.
3. The system of claim 1 , wherein the memory includes executable instructions that when executed by the processor perform actions that:
execute a first beam search that uses the non-terminal selector model to predict the select non-terminal symbol to expand; and
execute a second beam search that uses the non-terminal expansion model to predict the select production rule to expand the select non-terminal symbol.
4. The system of claim 1 , wherein the code sketch generates source code to complete a line of source code, a method body, a class, an expression, or a consecutive set of lines of source code.
5. The system of claim 1 , wherein the non-terminal selector model is a multilayer perceptron model trained through reinforcement learning.
6. The system of claim 1 , wherein the non-terminal expansion model is an encoder-decoder neural transformer model with attention trained through reinforcement learning.
7. A computer-implemented method for grammar-based code generation, comprising:
accessing a partially-formed source code snippet of a source code program associated with a programming language, wherein the programming language is associated with a grammar comprising a plurality of production rules, a plurality of non-terminal symbols, and a plurality of terminal symbols, wherein a production rule is a rule that replaces a first non-terminal symbol with one or more terminal symbols or one or more non-terminal symbols;
transforming the partially-formed source code snippet into a partial code state comprising at least one non-terminal symbol;
generating a code sketch to complete the partially-formed source code snippet, wherein the code sketch comprises a sequence of terminal symbols and a hole, wherein generating the code sketch comprises invoking a non-terminal selector model to predict a non-terminal symbol to expand in the partial code state and a non-terminal expansion model to predict the production rule that expands the predicted non-terminal symbol,
wherein the hole represents an unexpanded non-terminal symbol.
8. The computer-implemented method of claim 7 , further comprising:
outputting the code sketch in the source code program.
9. The computer-implemented method of claim 7 , further comprising:
executing a first beam search that invokes the non-terminal selector model to predict top-m non-terminal symbols to expand; and
for each of the top-m predicted non-terminal symbols, executing a second beam search that invokes the non-terminal expansion model to predict top-n expansions of each of the top-m non-terminal symbols.
10. The computer-implemented method of claim 9 , wherein the non-terminal selector model generates an output probability for each non-terminal symbol in the plurality of non-terminal symbols.
11. The computer-implemented method of claim 10 , wherein the top-m predicted non-terminal symbols have highest output probabilities.
12. The computer-implemented method of claim 9 , wherein the non-terminal expansion model generates an output probability for each production rule.
13. The computer-implemented method of claim 12 , wherein the top-n expansions have highest output probabilities.
14. The computer-implemented method of claim 9 , further comprising:
expanding each of the predicted top-m non-terminal symbols until a special stop symbol is predicted.
15. The computer-implemented method of claim 9 , further comprising:
terminating expansion of a select one of the predicted top-m non-terminal symbols when the non-terminal expansion model predicts a hole to replace a non-terminal symbol.
16. A hardware storage device having stored thereon computer executable instructions that are structured to be executable by a processor of a computing device to thereby cause the computing device to perform actions for grammar-based code completion that:
obtain a partially-formed source code snippet of a source code program in a programming language associated with a grammar comprising a plurality of production rules, a set of non-terminal symbols, and a set of terminal symbols, wherein a production rule is a rule that replaces a non-terminal symbol with one or more terminal symbols or one or more non-terminal symbols;
transform the partially-formed source code snippet into a partial code state comprising at least one non-terminal symbol and a squash non-terminal symbol, wherein the squash non-terminal symbol represents a position where a code sketch is to be inserted into the partially-formed source code snippet; and
generate a code sketch to insert into the position of the squash non-terminal symbol, wherein the code sketch comprises a sequence of terminal symbols and a hole, wherein the generation of the code sketch comprises invoking a non-terminal selector model to predict a non-terminal symbol in the partial code state to expand and a non-terminal expansion model to predict the production rule that expands the predicted non-terminal symbol,
wherein the hole represents an unexpanded non-terminal symbol.
17. The hardware device of claim 16 , wherein the partial code state comprises a first sequence of tokens, the squash non-terminal symbol, and a second sequence of tokens, wherein the first sequence of tokens represents source code preceding the code sketch, wherein the second sequence of tokens represents source code following the code sketch.
18. The hardware device of claim 16 , having stored thereon computer executable instructions that are structured to be executable by the processor of the computing device to thereby cause the computing device to perform actions for grammar-based code completion that:
output the code sketch in the source code program.
19. The hardware device of claim 16 , wherein the non-terminal selector model is a multi-layer perceptron model.
20. The hardware device of claim 16 , wherein the non-terminal expansion model is a neural transformer model with attention.