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:
receiving a training dataset of unlabeled code corpus comprising at least a PL segment and a natural language (NL) segment;
generating a bimodal input sequence by concatenating the PL segment and the NL segment;
generating, by the encoder-decoder model, a first predicted output in response to the bimodal input sequence according to a first pre-training task,
wherein generating the first predicted output comprises masking a plurality of identifiers in the PL segment of the bimodal input sequence;
computing a first training objective based on the first predicted output according to the first pre-training task;
generating, by the encoder-decoder model, a second predicted output in response to the bimodal input sequence according to a second pre-training task;
wherein the second pre-training task is different from the first pre-training task;
computing a second training objective based on the second predicted output according to the second pre-training task; and
alternately updating the encoder-decoder model based on the first training objective or the second training objective according to a pre-defined probability for each pre-training task.
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; encoding, by an encoder, the masked bimodal input sequence into a first representation; and generating, by a 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 a designated mask token is used for a specific identifier; encoding, by an encoder, the masked bimodal input sequence into a first representation; and generating, by a decoder, a target sequence comprising the masked plurality of identifiers and corresponding designated mask tokens.
5. The method of claim 4 , 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.
6. The method of claim 1 , further comprising:
converting the PL segment into an Abstract Syntax Tree (AST); and
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.
7. The method of claim 6 , wherein the first predicted output is generated by:
encoding, by an encoder, the bimodal input sequence into a first representation; and
mapping final hidden states corresponding to the PL segment in the encoder into a sequence of probabilities.
8. The method of claim 7 , 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.
9. The method of claim 1 , further comprising:
inputting the PL segment to the encoder-decoder model to generate a predicted NL segment;
computing a first loss by comparing the predicted NL segment and the NL segment;
inputting the NL segment to the encoder-decoder model to generate a predicted PL segment;
computing a second loss by comparing the predicted PL segment and the PL segment; and
updating the encoder-decoder model based on a weighted sum of the first loss and the second loss via backpropagation.
10. The method of claim 1 , further comprising:
receiving one or more training datasets for training the encoder-decoder 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 encoder-decoder model, a predicted output in response to the input training sequence according to the task control code.
11. The method of claim 10 , 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.
12. A system for programming language (PL) generation and understanding using an encoder-decoder model, the system comprising:
a communication interface receiving a training dataset of unlabeled code corpus comprising at least a PL segment and a natural language (NL) segment;
a memory storing a plurality of processor-executable instructions; and
a processor executing the plurality of processor-executable instructions to perform operations comprising:
generating a bimodal input sequence by concatenating the PL segment and the NL segment;
generating, by the encoder-decoder model, a first predicted output in response to the bimodal input sequence according to a first pre-training task,
wherein generating the first predicted output comprises masking a plurality of identifiers in the PL segment of the bimodal input sequence;
computing a first training objective based on the first predicted output according to the first pre-training task;
generating, by the encoder-decoder model, a second predicted output in response to the bimodal input sequence according to a second pre-training task;
wherein the second pre-training task is different from the first pre-training task;
computing a second training objective based on the second predicted output according to the second pre-training task; and
alternately updating the encoder-decoder model based on the first training objective or the second training objective according to a pre-defined probability for each pre-training task.
13. The system of claim 12 , wherein the first predicted output is generated by: randomly masking a plurality of spans in the bimodal input sequence; encoding, by an encoder, the masked bimodal input sequence into a first representation; 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, and 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.
14. The method system of claim 12 , wherein a designated mask token is used fora specific identifier; encoding, by an encoder, the masked bimodal input sequence into a first representation; and generating, by a decoder, a target sequence comprising the masked plurality of identifiers and corresponding designated mask tokens, 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.
15. The system of claim 12 , wherein the operations further comprises:
converting the PL segment into an Abstract Syntax Tree (AST); and
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
wherein the first predicted output is generated by:
encoding, by an encoder, the bimodal input sequence into a first representation; and
mapping final hidden states corresponding to the PL segment in the encoder into a sequence of probabilities,
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. 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 instructions being executed by a processor to perform operations comprising:
receiving a training dataset of unlabeled code corpus comprising at least a PL segment and a natural language (NL) segment;
generating a bimodal input sequence by concatenating the PL segment and the NL segment;
generating, by the encoder-decoder model, a first predicted output in response to the bimodal input sequence according to a first pre-training task,
wherein generating the first predicted output comprises masking a plurality of identifiers in the PL segment of the bimodal input sequence;
computing a first training objective based on the first predicted output according to the first pre-training task;
generating, by the encoder-decoder model, a second predicted output in response to the bimodal input sequence according to a second pre-training task;
wherein the second pre-training task is different from the first pre-training task;
computing a second training objective based on the second predicted output according to the second pre-training task; and
alternately updating the encoder-decoder model based on the first training objective or the second training objective according to a pre-defined probability for each pre-training task.
17. The non-transitory machine-readable medium of claim 16 , wherein the first predicted output is generated by: randomly masking a plurality of spans in the bimodal input sequence; encoding, by an encoder, the masked bimodal input sequence into a first representation; and generating, by a 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, and 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.
18. The non-transitory machine-readable medium of claim 16 , wherein a designated mask token is used for a specific identifier; encoding, by an encoder, the masked bimodal input sequence into a first representation; and generating, by a decoder, a target sequence comprising the masked plurality of identifiers and corresponding designated mask tokens, 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.
19. The non-transitory machine-readable medium of claim 16 , wherein the operations further comprise: converting the PL segment into an Abstract Syntax Tree (AST); and 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 wherein the first predicted output is generated by: encoding, by an encoder, the bimodal input sequence into a first representation; and mapping final hidden states corresponding to the PL segment in the encoder into a sequence of probabilities, 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.