Learning to Complete Code with Sketches

1 INTRODUCTION . Recent high-capacity language models ( LM ) have shown that machine learning models are able to generate coherent , realistic text , but it is often hard to guide them towards a specific goal , especially when describing the intent is complex or more costly than manually generating the target output . One such scenario are LMs of source code ( LMC ) . Since Hindle et al . ( 2012 ) increasingly sophisticated LMCs have been built , including transformer-based ones , such as those of Svyatkovskiy et al . ( 2020 ) ; Feng et al . ( 2020 ) ; Chen et al . ( 2021 ) and various similar unpublished models such as TabNine and SourceAI . These models generate full sequences of code tokens left-to-right with any prefix acting as the ( partial ) user intent . While LMs generate realistic-looking outputs , they are known to occasionally “ hallucinate ” ( Puduppully et al. , 2019 ; Malmi et al. , 2019 ; Maynez et al. , 2020 ; Liu et al. , 2021 ) , i.e . generate plausible but incorrect content . This is particularly problematic in generating source code , where small mistakes can lead to erroneous code that is very hard to debug or introduces vulnerabilities ( Pearce et al. , 2021 ) . In this work , we investigate models that can decline to make predictions in places where there is high uncertainty ( e.g. , where the user should choose a name ) , but continue generating around these “ holes ” . For example , in Fig . 1 ( left ) a developer has typed some code and is about to type the next line . A likely completion is to consume more command line arguments , but their name is unclear Code Context : 1 import sys 2 target = sys.argv [ 1 ] 3 I Ground-Truth : ID = sys.argv [ 2 ] Suggested Code Completions : L→ R target = target.replace ( `` \\ '' , `` / '' ) L→ R +⦸ target = L→ R + print ( target ) Copilot ( No suggestion ) GRAMMFORMER = sys.argv [ 2 ] Figure 1 : A sample snippet ( left ; abbreviated from Fig . 12 in Appx . A ) . A developer has just typed the code and their cursor ( in blue ) is at line 3 . Code completions provided by a number of models are shown on the right , where L→ R is a standard LMC and GRAMMFORMER is our new model . from the context . A traditional generative model ( e.g . Fig . 1 ; top right ) may choose to provide a completion that exists in the training data , but is not clearly called for here . On the other hand , a model able to explicitly mark where it is uncertain ( Fig . 1 ; bottom right ) makes it clear to a user where further input is required . However , creating such models is not trivial . A simple first attempt may be to use a standard LMC , but output a “ hole token ” whenever the model is uncertain about the next output token . However , continuing after the “ ” then becomes infeasible , as the LMC was not trained on such data . Hence , a suitable training dataset and objective need to devised . As no large datasets with holes exist , we instead choose to use a reinforcement learning approach in which our reward function encourages the model to make “ long ” predictions with as few “ ” tokens as possible , but to avoid making incorrect predictions . We found that standard left-to-right sequence models perform poorly on this task . Hence , we developed GRAMMFORMER , a model that construct suggestions by generating a ( partial ) syntax tree , but which has the option of leaving non-terminals in its output . Contributions ( 1 ) We present GRAMMFORMER , a transformer-based model that generates code based on the programming language grammar and can predict hole tokens rather than output it is uncertain about . ( 2 ) We develop REGEXACC , a metric that evaluates the quality of predictions with holes . ( 3 ) We evaluate GRAMMFORMER on Python and C # code and show that GRAMMFORMER makes longer and more precise statement-level sketch completions compared to baselines . 2 METHOD . Our aim is to predict code completions as sketches , a mix of actual tokens and “ holes ” , which are meant to signify that the model is unable to make a useful prediction within the given context and further user input is required . Formally , we consider models that take a context sequence x of tokens as input and have to produce an output sequence y ; intuitively , x is what the user typed so far , and y is the suggestion presented to the user . In our setting , y is a sketch , a mix of tokens from the programming language and the special token signifying a “ hole ” that could be filled by an arbitrary sequence of tokens . For example , t = foo ( ) is a sketch corresponding to assigning the return value of function foo to variable t , but leaves the arguments of the function call undetermined . Metric A good sketch is one that ( a ) can be completed into the correct output and ( b ) is as precise as possible . To measure how successful a method is in doing so , we define a new metric REGEXACC . For ( a ) , we use toRegex ( ŷ ) to turn a predicted code sketch ŷ into a regular expression by replacing all holes with the wildcard matching any non-empty sequence ( “ .+ ” in Perl Compatible Regular Expression syntax ) . If the regex matches the ground truth , matches ( ⋅ , ⋅ ) returns a score of 1 otherwise it returns 0 . To implement ( b ) , we scale this result by the proportion of terminal tokens predicted , by defining nTokens ( ŷ ) as the function that returns the number of non-hole symbols in ŷ . More formally , assume an output sketch ŷ and a ground-truth sequence y∗ , where y∗ does not contain any tokens . REGEXACC is then defined as REGEXACC ( ŷ , y∗ ) ≜ matches ( toRegex ( ŷ ) , y∗ ) ⋅ nTokens ( ŷ ) nTokens ( y∗ ) . Beyond REGEXACC , we also consider ROUGE ( Lin , 2004 ) , since a sketch can be thought as a form of a “ summary ” of the target text . For this , we use a helper function ERASEHOLES ( ŷ ) that simply drops all tokens , and then consider ROUGEF1 ( ERASEHOLES ( ŷ ) , y∗ ) . ROUGE is more lenient to errors than REGEXACC and gives partial credit to non-matching but plausible sketches . 2.1 LINEAR CODE SKETCH GENERATION . First , we consider the idea of generating code sketches using a standard generative model for language . To this end , we simply extend the vocabulary with the special “ ” token . An obvious problem is that while we have plenty of training data for a standard generative model , we do not have training data for outputs y that contain the token . Consequently , we can not train the model in a fully supervised fashion , and instead turn to reinforcement learning . Concretely , we devise a reward function r ( ⋅ ) that averages REGEXACC and ROUGE , i.e . for a predicted output sketch ŷ and a ground truth output ( without tokens ) y∗ , we define r ( ŷ , y∗ ) = 1 2 ( REGEXACC ( ŷ , y∗ ) + ROUGEF1 ( ERASEHOLES ( ŷ , y∗ ) ) . ( 1 ) Using the combination of ROUGE ( which does not consider holes ) and REGEXACC is crucial here , as ROUGE is much “ smoother ” compared to REGEXACC , which is 0 for all but very few predictions , allowing us to measure partial improvement . We use our reward function from Eq . 1 to evaluate the quality of the output of the full model and compute a loss . Inspired by Paulus et al . ( 2017 ) we use self-critical policy gradient training ( Rennie et al. , 2017 ) and for a prediction ŷ we minimise L ( x , y∗ ) = ( r ( ŷ , y∗ ) − r̃ ( x ) ) ⋅ Lgen ( x , ŷ ) ( 2 ) Here , r̃ ( x ) is the reward achieved by the prediction from the snapshots of the model that achieved the best score so far and Lgen is the loss of the generative model . Intuitively , this objective rewards models that improve upon the previous best policy with respect to r. To model this in practice , we use a standard encoder/decoder Transformer model Vaswani et al . ( 2017 ) ; Radford et al . ( 2019 ) , “ translating ” the context x into the output y using separate encoder and decoder models . We additionally also consider the language modelling case , i.e. , a model that conditioned on x predicts token y0 , conditioned on x , y0 predicts token y1 , etc .. Pretraining In practice , we found that directly training a sequence model to maximise Eq . 1 , is very slow and does not converge to a useful model . Instead , we heuristically generate a dataset suitable for supervised pretraining . We replace random AST non-terminals of the target output by and generate target sequences . These contain terminals and zero or more . We then pretrain the model on this dataset to convergence , and then fine-tune it using the reward of Eq . 1 . 2.2 GRAMMAR-BASED CODE SKETCH GENERATION . In experiments , we found the simple extended sequence model from above to not perform well , in particular , tokens would not replace semantically meaningful subsequences ( e.g . “ szconv . ) ” Algorithm 1 GRAMMFORMER generative process , given an input sequence x ( 0 ) . for t = 0 , 1 , 2 , ... do i ( t ) ∼ Ps ( i ∣ x ( t ) , N ( x ( t ) ) ) ▷ sample non-terminal position from N ( x ( t ) ) to expand if i ( t ) = ⦸ then ▷ if x ( t ) does not contain non-terminals or none was selected by Ps break ▷ stop generation u ( t ) ⊚i ( t ) ∼ Pe ( u ∣ x ( t ) , i ( t ) ) ▷ sample expansion of non-terminal at position i ( t ) x ( t+1 ) ← x ( t ) < i ( t ) ∶∶ u ( t ) ⊚i ( t ) ∶∶ x ( t ) > i ( t ) ▷ create x ( t+1 ) by replacing non-terminal at i ( t ) by u ( t ) ⊚i ( t ) return NONTERMINALSTOHOLES ( x ( t ) ) ▷ convert remaining non-terminals to holes and return does not contain a left parenthesis and requires the user to fill it in. ) . To resolve this , we developed GRAMMFORMER , a grammar-guided model . It generates code by following the structure of the context-free grammar ( CFG ) defining the programming language syntax , iteratively expanding nonterminal symbols . Crucially , it can choose to not expand some non-terminal symbols , which can then be presented as to users . In traditional grammar-based generation of text ( Cohen et al. , 2012 ) or code ( Maddison & Tarlow , 2014 ; Yin & Neubig , 2017 ; Allamanis & Sutton , 2014 ; Bielik et al. , 2016 ) , the CFG is followed by sequentially expanding the left-most , bottom-most non-terminal symbol , using one of the production rules of the grammar . GRAMMFORMER changes this and instead selects the non-terminal symbol to expand , if any . An example generation is shown in Fig . 2 . Probabilistic Model A CFG is defined as a tuple ( Σ , N , S , R ) where Σ is a set of terminal symbols , N is a set of non-terminal symbols , S ∈ N is the root symbol andR is a set of production rules . We denote non-terminals as ⟨NonTerminalName⟩ . GRAMMFORMER can be viewed as a sequenceto-sequence model transforming x = x0 , x1 , ... , xn into a new sequence in which one non-terminal symbol xi has been replaced by a new sequence of new symbols , according to a production rule of the grammar . Examples of such sequences and rewrites are shown in Fig . 2 . GRAMMFORMER does this rewriting in two steps . First , a non-terminal selector model Ps selects a non-terminal in x to expand and then the non-terminal expansion model Pe determines how to expand it . To define Ps , letN ( x ) = { i ∣ xi ∈ N } ∪ { ⦸ } denote the set of non-terminal positions in x and a special “ stop expansion ” ⦸ symbol . Conditioned on x , Ps produces a probability distribution over N ( x ) . In turn , Pe is conditioned on x and a position i ∈ N ( x ) and models a probability distribution over expansion sequences u ∈ ( Σ ∪ N ) ∗ . Note that factorising GRAMMFORMER into two models Ps and Pe is an important modelling decision : how to best expand a non-terminal is entirely separated from predicting whether a hole should be introduced . These two concepts are intermixed in standard ( sequence ) decoders . In practice , we define both models using neural architectures with partially shared parameters , as discussed below . Alg . 1 shows a high-level description of GRAMMFORMER , in which Ps and Pe are used repeatedly to select and expand non-terminals ( not necessarily the left-most one ) , until none are left or the Ps indicates that expansion should stop . Here , NONTERMINALSTOHOLES ( ⋅ ) replaces all remaining non-terminal symbols with a hole . Note that GRAMMFORMER is not context-free , taking into account the whole input sequence when expanding a non-terminal . Second , in contrast to many grammar-based methods ( Yin & Neubig , 2017 ; Bielik et al. , 2016 ) , any non-terminal can be expanded at each step . Finally , Pe is not directly constrained to follow the production rule set R , but can generate any sequence . In practice , it learns to follow to the rules of R from the data , but this flexibility is important for handling string literals and argument tuples of variable length . Neural Model To implement Ps and Pe , we use a shared encoder module that computes a representation of the input sequence x = x0 , . . . , xn as vectors e0 , . . . , en , ei ∈ R D , where D is a hyperparameter . Our encoder module is a Transformer ( Vaswani et al. , 2017 ) , given the impressive results of transformer-based models in NLP and code ( Feng et al. , 2020 ) . Other architectures ( RNNs , 1D-CNNs , Transformer variants ) would be suitable , but we leave their study for future work . Ps is implemented similar to a pointer network on top of this encoder module , i.e . Ps ( i ∣ x ) = softmax i∈N ( x ) ( f ( ei ) ) , where f is a learnable feed-forward neural network . For our purposes , we define e⦸ as the representation of the special start symbol [ CLS ] used in our Transformer encoder . The expansion model Pe follows a standard autoregressive decoder formulation , i.e . Pe ( u ∣ x , i ) = m ∏ j=1 Pdec ( uj ∣ e0 , . . . , en , i , u < j ) . We implement Pdec as a ( causal ) relational Transformer decoder , similar to Wang et al . ( 2019 ) . Relational transformers augment the attention mechanism by incorporating predefined relationships among elements ; attention scores are then biased by learnable weights for each relation . In GRAMMFORMER , we only use a single relation , connecting each token to the expanded non-terminal token xi , to help the model focus on the token it needs to generate an expansion for . Objective Due to the lack of supervised data , we employ reinforcement learning to train GRAMMFORMER . We use our reward function from Eq . 1 to evaluate the quality of the output of the full model . We use self-critical policy gradient training as in Eq . 2 and minimise L ( x , y∗ ) = ( r ( ŷ , y∗ ) − r̃ ( x ) ) ⋅ T ∑ t=0 ( − logPs ( i ( t ) ∣ x ( t ) ) − I ( i ( t ) ≠ ⦸ ) ⋅ logPe ( ( u ( t ) ⊚i ( t ) ) ∗ ∣ x ( t ) , i ( t ) ) ) . ( 3 ) Here , r̃ ( x ) is the reward achieved by the snapshots of Ps and Pe that achieved the best score so far . The rest of the objective follows the iterations of the loop in Alg . 1 , where t is the iteration index , ŷ is the predicted sketch , y∗ is the ground-truth sequence of terminals , and I ( ⋅ ) is the indicator function . Pretraining As in the sequence model , directly training with the RL objective Eq . 3 is computationally intensive due to the sampling requirement . We again use a pretraining strategy . First , we train Pe to expand every non-terminal , independently of the expansion order learned by Ps . To do this , we use the input training examples and follow Alg . 1 , but instead of sampling from Ps ( ⋅ ) , we sample i ( t ) from a uniform distribution over the non-terminals in x ( t ) , Ñ ( x ( t ) ) = { i ∣ xi ∈ N } . This yields sequences of intermediate sketches x ( t ) for each example . Furthermore , for each x ( t ) , we compute the ground-truth expansion ( u ( t ) ⊚i ) ∗ for all non-terminals i ∈ Ñ ( x ( t ) ) . We can then pretrain Pe using the supervised objective Lpre , e ( x ( t ) , ( u ( t ) ⊚i ) ∗ i∈Ñ ( x ( t ) ) ) = 1 ∣Ñ ( x ( t ) ) ∣ ⋅ ∑ i∈Ñ ( x ( t ) ) − logPe ( ( u ( t ) ⊚i ( t ) ) ∗ ∣ x ( t ) , i ) , i.e . the negative log-likelihood of the correct expansion for all non-terminals in x ( t ) . This computation is more computationally efficient compared to the one in Eq . 3 since the cost of encoding x ( t ) is amortised across all potential expansions and no sampling is required . Once Pe is pretrained , we pretrain Ps . For this , we fix the weights of the shared encoder module , and optimise only the remaining parameters of Ps through Eq . 3 . Once we have a pretrained both models , we then fine-tune all model weights end-to-end , using Eq . 3 . Optimisation : Grammar Flattening Following the formal grammar of a programming language commonly introduces tedious expansions . For example , the Python non-terminal ⟨Call⟩ is always expanded to ⟨Expr⟩ ( ⟨ArgumentList⟩ ) , and the C # non-terminal ⟨NotEqualOp⟩ is always expanded to the terminal ! = . We “ flatten ” the grammar by replacing non-terminals such as ⟨Call⟩ and ⟨NotEqualOp⟩ with all their possible expansions . In Appx . C we provide the list of the flattened non-terminals . Note that if we repeated this process for all non-terminals except from the starting symbol S , GRAMMFORMER would degenerate into a standard encoder-decoder model . Beam Search At test time , we employ a two-step beam search , and replace sampling from Ps and Pe with their top-ν outputs , keeping a beam of size k. First , for each x ( t ) in the beam , we compute Ps and select the top-m non-terminal positions to expand . For each of those m positions , we sample the top-n expansions from Pe using a standard beam search . We compute the likelihood of all k ⋅ n ⋅m results , and then keep only the top-k . This process ( detailed in Appx . E ) is similar to a standard beam search but takes into account that two submodels are used . Computational Cost GRAMMFORMER ’ s ability to predict sketches comes with additional computational cost compared to standard transformer encoder-decoders : at each iteration of the loop in Alg . 1 x ( t ) changes , Ps and Pe must be recomputed . This means that the encoder-decoder runs once on each partial sequence , in contrast to left-to-right causal generation , in which intermediate results can be re-used . Future work may consider selecting more than one element to expand from N ( x ( t ) ) at each step , reducing the expansion steps , similar to Welleck et al . ( 2019 ) ; Stern et al . ( 2019 ) .

Summary: The paper presents a new model for code completion which allows the model to completions with “holes” that are inserted in places where the model is uncertain. The idea of generating “holes” to enable skipping over “hard parts” of the prediction is novel and interesting. To realize this idea, the authors present a model that generates partial syntax trees where some of the non-terminals may be left without further expansion.  The model is evaluated on C# and Python programs and is shown to outperform existing techniques. The also paper presents a thorough ablation study.