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: This work proposes, GRAMMAFORMER, a transformer model for generating code with "holes" inserted in places where a model is uncertain. GRAMMAFORMER is trained on code completion task for C# and Python. The model generates 10-50% more accurate completions and 37-50% longer sketches.