Systems and methods for code understanding and generation
Embodiments described herein a code generation and understanding model that builds on a Transformer-based encoder-decoder framework. The code generation and understanding model is configured to derive generic representations for programming language (PL) and natural language (NL) in code domain via pre-training on unlabeled code corpus, and then to benefit many code-related downstream tasks with fine-tuning. Apart from the denoising sequence-to-sequence objectives widely adopted for pre-training on natural language, identifier tagging and prediction pre-training objective is adopted to enable the model to better leverage the crucial token type information from PL, which specifically are the identifiers assigned by developers.
1. A method for programming language (PL) generation and understanding using an encoder-decoder model, the method comprising:
generating, by a neural network model, a first predicted output in response to a bimodal input sequence comprising a PL segment and a natural language (NL) segment according to a first pre-training task, wherein the first predicted output is generated by:
masking a plurality of identifiers in the PL segment of the bimodal input sequence to generate a masked bimodal input sequence, wherein a designated mask token is used for a specific identifier of the plurality of identifiers,
encoding, by an encoder of the neural network model, the masked bimodal input sequence into a first representation, and
generating, by a decoder of the neural network model, a target sequence comprising the masked plurality of identifiers and corresponding designated mask tokens, from the first representation to generate the first predicted output
wherein the PL segment and the NL segment belong to an unlabeled code sample;
computing a first training objective based on the first predicted output according to the first pre-training task;
generating, by the neural network model, a second predicted output in response to the bimodal input sequence according to a second pre-training task;
computing a second training objective based on the second predicted output according to the second pre-training task; and
alternately or jointly updating the neural network model based on the first training objective and the second training objective.
2. The method of claim 1 , wherein the first predicted output is generated by: randomly masking a plurality of spans in the bimodal input sequence; and generating, by the decoder, a predicted masked sequence conditioned on the masked bimodal input sequence and a set of previously predicted masked sequences up to a current time instance.
3. The method of claim 2 , wherein the first training objective is computed based on a conditional probability distribution of the masked bimodal input sequence and the set of previously predicted masked sequences up to the current time instance.
4. The method of claim 1 wherein the first training objective is computed based on a conditional probability of each token in the target sequence conditioned on the masked bimodal input sequence and tokens previously generated for the target sequence up to a respective token.
5. The method of claim 1 , wherein the first predicted output is generated by: converting the PL segment into an Abstract Syntax Tree (AST); constructing, based on the AST, a sequence of binary labels, each binary label representing whether a respective code token in the PL segment is an identifier; and mapping final hidden states corresponding to the PL segment in the encoder into a sequence of probabilities.
6. The method of claim 5 , wherein the first training objective is computed based on a binary cross-entropy loss between the sequence of probabilities and the sequence of binary labels.
7. The method of claim 1 , wherein the second pre-training task is different from the first pre-training task.
8. The method of claim 1 , further comprising: receiving one or more training datasets for refining the neural network model on multiple tasks at a time; prepending an input training sequence with a task control code indicating a type of a task corresponding to a specific training dataset; and generating, by the neural network model, a predicted output in response to the input training sequence according to the task control code.
9. The method of claim 8 , wherein the input training sequence is sampled from the one or more training datasets according to a multinomial distribution computed based on respective sizes of the one or more training datasets.
10. A system for programming language (PL) generation and understanding using a neural network model, the system comprising:
a memory storing the neural network model comprising an encoder and a decoder, and a plurality of processor-executable instructions; and
a processor executing the plurality of processor-executable instructions to perform operations comprising:
generating, by the neural network model, a first predicted output in response to a bimodal input sequence comprising a PL segment and a natural language (NL) segment according to a first pre-training task, wherein the first predicted output is generated by:
masking a plurality of identifiers in the PL segment of the bimodal input sequence to generate a masked bimodal input sequence, wherein a designated mask token is used for a specific identifier of the plurality of identifiers,
encoding, by an encoder of the neural network model, the masked bimodal input sequence into a first representation, and
generating, by a decoder of the neural network model, a target sequence comprising the masked plurality of identifiers and corresponding designated mask tokens, from the first representation to generate the first predicted output,
wherein the PL segment and the NL segment belong to an unlabeled code sample;
computing a first training objective based on the first predicted output according to the first pre-training task;
generating, by the neural network model, a second predicted output in response to the bimodal input sequence according to a second pre-training task;
computing a second training objective based on the second predicted output according to the second pre-training task; and
alternately or jointly updating the neural network model based on the first training objective and the second training objective.
11. The system of claim 10 , wherein the first predicted output is generated by: randomly masking a plurality of spans in the bimodal input sequence; and generating, by the decoder, a predicted masked sequence conditioned on the masked bimodal input sequence and a set of previously predicted masked sequences up to a current time instance.
12. The system of claim 11 , wherein the first training objective is computed based on a conditional probability distribution of the masked bimodal input sequence and the set of previously predicted masked sequences up to the current time instance.
13. The system of claim 10 , wherein the first training objective is computed based on a conditional probability of each token in the target sequence conditioned on the masked bimodal input sequence and tokens previously generated for the target sequence up to a respective token.
14. The system of claim 10 , wherein the first predicted output is generated by: converting the PL segment into an Abstract Syntax Tree (AST); constructing, based on the AST, a sequence of binary labels, each binary label representing whether a respective code token in the PL segment is an identifier; and mapping final hidden states corresponding to the PL segment in the encoder into a sequence of probabilities.
15. The system of claim 14 , wherein the first training objective is computed based on a binary cross-entropy loss between the sequence of probabilities and the sequence of binary labels.
16. The system of claim 10 , wherein the second pre-training task is different from the first pre-training task.
17. The system of claim 10 , wherein the operations further comprise: receiving one or more training datasets for refining the neural network model on multiple tasks at a time; prepending an input training sequence with a task control code indicating a type of a task corresponding to a specific training dataset; and generating, by the neural network model, a predicted output in response to the input training sequence according to the task control code.
18. A non-transitory machine-readable medium storing a plurality of processor-executable instructions for programming language (PL) generation and understanding using an encoder-decoder model, the plurality of processor-executable instructions being executed by one or more hardware processors to perform operations comprising:
generating, by the neural network model, a first predicted output in response to a bimodal input sequence comprising a PL segment and a natural language (NL) segment according to a first pre-training task, wherein the first predicted output is generated by:
masking a plurality of identifiers in the PL segment of the bimodal input sequence to generate a masked bimodal input sequence, wherein a designated mask token is used for a specific identifier of the plurality of identifiers,
encoding, by an encoder of the neural network model, the masked bimodal input sequence into a first representation, and
generating, by a decoder of the neural network model, a target sequence comprising the masked plurality of identifiers and corresponding designated mask tokens, from the first representation to generate the first predicted output,
wherein the PL segment and the NL segment belong to an unlabeled code sample;
computing a first training objective based on the first predicted output according to the first pre-training task;
generating, by the neural network model, a second predicted output in response to the bimodal input sequence according to a second pre-training task;
computing a second training objective based on the second predicted output according to the second pre-training task; and
alternately or jointly updating the neural network model based on the first training objective and the second training objective.