CodePAD: Sequence-based Code Generation with Pushdown Automaton

Code generation is a hot research spot in the field of natural language processing (NLP) and software engineering Ling et al. (2016); Yin and Neubig (2018); Wei et al. (2019); Dong et al. (2022); Shen et al. (2022), which aims to help facilitate software development and revolutionize user programming. A well-designed code generation approach requires consideration of grammar in addition to semantics, as semantics ensures the functionality of code meets developer’s intention, while satisfying grammar is a basic requirement for programming.

Sequence-based approaches are the most commonly used in code generation Raychev et al. (2014); Jia and Liang (2016); Cao et al. (2019); Wang et al. (2021); Fried et al. (2022); Nijkamp et al. (2022), which has the following three advantages: 1) High efficiency. Sequence-based approach directly generates a token sequence of code along the order of human writing in one pass. 2) Convenient operation. It can obtain partially generated code with ease even if the result is incomplete or the generation is not finished, thus can be adopted for various code generation scenarios ¹¹1For example, code completion, a popular variant of code generation in practice, which completes code with contexts.. 3) Easy data accessibility. Sequence structured data is the most general form that can be obtained without much effort compared to other structured data. Thus, pre-trained methods with big data-driven usually employ sequence-based approaches Wang et al. (2021); Ahmad et al. (2021); Guo et al. (2022); Fried et al. (2022); Nijkamp et al. (2022). However, most sequence-based approaches generate codes with no guarantee of grammatical correctness (GC), which hinders the usability of generated codes.

According to automata theory, we find that PDA is suitable for PL grammar and has the following properties for dealing with non-grammatical problems. The first property is a valid terminal symbol set for next token prediction, and the second is an accept state set for recognizing completed code. In a practical code generation process, what we really care about are two typical grammatical error: terminal symbol mismatch (TSM) error and end with non-accept state (ENS) error, where ENS error can be tolerable until the end of code generation process. As a result, satisfying grammar constraints during code generation require the above two properties, and constraining code generation based on PDA to guarantee GC is a matter of course.

It is inspired by the operating principle of programming language parsers in practice. Fig. 1 demonstrates some examples of a parser parsing python code, which contains two specific grammatical error codes (a) and (b), and a grammatically correct code (c). Case (a) indicates ENS error, and case (b) indicates TSM error. All cases are recognized by a parser according to grammar. Therefore, a straightforward way is applying a parser to constrain sequence-based code generation, but it has the following two problems. First, sequence-based models need to recognize grammatical errors during code generation process. However, the vast majority of codes generated during this situation are incomplete, i.e., parsing these codes will produce ENS errors like case (a). Second, a parser can only recognize TSM errors like case (b) if terminal symbol is fed to the parser. It implies that sequence-based models need to attempt each generated code token in the candidate set one by one, which is strenuous ²²2For example, there are only 83 syntax-strings and 10 token-types in Python, and for a 10k-length vocabulary, more than 99% of tokens belong to three token-types NAME, STRING, and NUMBER. If models want to generate these types of tokens based on semantics but are prevented by grammar, the cost of attempts could be overwhelming.. In short, although code can be grammatically bounded by a parser, it is still challenging to work directly with a parser on sequence-based code generation.

In this paper, we propose CodePAD (a Code generation framework based on Pushdown Automaton moDule) to ensure GC for sequence-based models. PDA module facilitates CodePAD to generate bounded next prediction in a valid set and end with an accept state, utilizing the principle that PDA can recognize grammatical codes. CodePAD not only ensures GC but also maintains the advantages of commonly used sequence-based code generation. We conduct a series of experiments on four public benchmark datasets. Extensive experimental results and analyses verify the effectiveness and generality of CodePAD. The main contribution of this paper can be summarized as follows:

• •

  We devise a PDA-based methodology to guarantee grammatical correctness for code generation, which consists of a PDA module and an algorithm to simulate the deduction of PDA.

• •

  We propose CodePAD, a novel sequence-based code generation framework incorporating the PDA module. In addition to using the PDA-based methodology, CodePAD leverages PDA state through state representation, state prediction task, and joint prediction with state, which helps learn PDA deduction.

• •

  CodePAD significantly enhances the performance of base models. Even in zero-shot setting, pre-trained models still show remarkable improvements.

An example of code generation constrained by PDA module is shown in Fig. 2. Given an NL, CodePAD needs to generate code that satisfies its intended functionality with the guarantee of GC. PDA module gives a valid set of terminal symbols (i.e., candidate syntax-strings and token-types) depending on grammar constraints, to which the generated token of CodePAD should belong. Specifically, in the beginning, CodePAD wants to generate first token from the vocabulary of a model according to NL. Based on production rules, PDA module shrinks the candidate set from entire vocabulary to 35 syntax-strings and 6 token-types, where each candidate token is constrained by grammar. As a result, at step 1, CodePAD picks a token (i.e., ‘import’) from the candidate set according to semantics. At each subsequent generation step, CodePAD predicts a token in the same manner, which satisfies grammar constraints. Until the end of generation steps, CodePAD outputs the token of ENDMARKER token-type (e.g., </s>), which ends with an accept state. Thus, CodePAD generates a grammatical code with the help of PDA module.

In the way of the above example, we first introduce PDA into deep learning (DL), which drastically shrinks the candidate set at each step of generation and reduces the difficulty for the model to select tokens from it, thus ensuring GC and improving the quality of code generation.

We first describe the core module of the proposed approach – a PDA-based methodology consisting of a PDA module and the corresponding algorithm to obtain a valid set for ensuring GC during code generation.

PLs belong to context-free languages, such as Python, Java, C++, etc., which can be accepted by a PDA Chomsky (1962); Schützenberger (1963); Evey (1963). A PDA ${\displaystyle M}$ can be defined as:

where $S$ is a finite set of states, $\Sigma$ is a finite set of input symbols, $\Gamma$ is a finite set of stack symbols, $s_{0}\in S$ is a start state, $g_{0}\in\Gamma$ is a starting stack symbol, $A\subseteq S$, where $A$ is the set of accept states, and $\delta$ is a transition function, mapping $S\,\times\Gamma\,\times(\Sigma\,\cup\left\{\varepsilon\,\right\})$ into the finite subsets of $S\times\Gamma^{*}$, where $\displaystyle*$ is Kleene star.

To construct a PDA module for a PL grammar, we first define $S$, $\Sigma$, $\Gamma$, $s_{0}$, $g_{0}$, and $A$. In this paper, we adopt non-terminal symbols $\mathcal{N}^{\prime}s$ of production rules in PL grammar as $S$, which can be changed with the PDA you build. $\Sigma$ is set to terminal symbols $\mathcal{T}^{\prime}s$ of production rules and $\Gamma$ is the union of $\mathcal{T}^{\prime}s$ and $\mathcal{N}^{\prime}s$. In PL grammar, $s_{0}$ and $g_{0}$ are the starting non-terminal symbol, and $A$ is the set of ending terminal symbols. As an example, in Python grammar PDA, $s_{0}$ and $g_{0}$ can be chosen as ‘file_input’, and $A$ is the set of ENDMARKER token-type. Then, the definition of $\delta$ is:

• •

  For each $\mathcal{N}$, $\delta$(s, $\mathcal{N}$,$\epsilon$) = {(s, $\beta$) | $\mathcal{N}\rightarrow\beta$ is a production rule in PL grammar}.

• •

  For each $\mathcal{T}$, $\delta$(s, $\mathcal{T}$, $\mathcal{T}$) = {(s, $\epsilon$)}.

The input symbols $\mathcal{I}^{\prime}s$ in $\Sigma$ are finite, because grammar of PLs merges the same type of tokens. For instance, each $\mathcal{I}$ in Python belongs to one of 83 syntax-strings and 10 token-types. As shown in Fig. 3, ‘import’ and ‘as’ are syntax-strings, while ‘numpy’ and ‘np’ belong to NAME token-type. Fig. 3 illustrates a Python grammar PDA parsing ‘import numpy as np’. For a valid $\mathcal{I}=$ ‘import’, PDA jumps to valid states and stacks based on $\delta$.

As shown in Algorithm 1, given current state $s$ and current stack $g$, PDA ${\displaystyle M}$ is able to provide the set of valid $\mathcal{I}^{\prime}s$ according to $\delta$. If $\varepsilon$ belongs to the set of valid $\mathcal{I}^{\prime}s$, the state and stack transferred after entering $\varepsilon$ should be taken into account as well. Eventually, we merge all sets of $\mathcal{I}^{\prime}s$ under each state and stack and remove $\varepsilon$ from them.

In this section, we introduce CodePAD, which is constrained by the PDA module constructed under methodology. This framework can be applied to any sequence-based model, no matter it is an encoder-decoder or a decoder only. In this paper, we employ Transformer-based encoder-decoder model Vaswani et al. (2017) as the backbone. We mainly modify the decode side, which adds PDA module, state representation, state prediction task, and joint prediction to help models generate grammatical code.

Fig. 4 shows the model architecture of CodePAD. At the encoder side of models, it maps an input sequence of NL utterances $\bm{x}=\left\{x_{1},x_{2},\cdots,x_{n}\right\}$ to a sequence of continuous NL representations $\bm{z}=\left\{z_{1},z_{2},\cdots,z_{n}\right\}$. At the decoder side of models, given the NL representation sequence $\bm{z}$, it then generates a token sequence of code, one element at a time. At each step, models generate the next token in an auto-regressive manner Graves (2013), which deploys previous generation results as additional input.

We conduct experiments to answer the following research questions (RQs):

There are three major limitations of our work:

• •

  First, a specific PDA should be built for each PL, and we only built PDA for Python (including both Python2 and Python3). However, PDA can recognize context-free language to which PL belongs, due to time and memory limitations of running PL Salomaa et al. (2001). In future work, we will build PDAs for more PLs.

• •

  Second, our work ensures GC, but neither our method nor other methods can prevent compilation errors. With our proposed training framework, our method outperforms other methods on compilation rates, e.g., CODEP has 99.2% compilation rates after training on CONALA.

• •

  Third, using PDA will increase CPU operation and memory overhead, but it is still acceptable.

In this paper, we have proposed a sequence-based framework, namely CodePAD, that first integrates the deduction of PDA into deep learning for code generation. CodePAD adds PDA module, state representation, state prediction task, and joint prediction for the decoder of CodePAD. PDA module guarantees the GC of generated codes, while the others make use of valid information provided by PDA module (i.e., PDA state) to generate codes better. As a result, CodePAD dramatically enhances the performance of base models and outperforms SOTA tree-based methods without pre-training on three benchmark datasets. The ablation study demonstrates each component of CodePAD contributes. With the help of PDA module, pre-trained models also achieve significant improvements.

PDA module shows excellent potential in code generation tasks in zero-shot setting, which can help pre-trained models apply to niche PLs that have little data. Furthermore, PDA can accommodate context-free language to satisfy grammar constraints, not just PL. We hope this work sheds light on future work in this direction.