Fine-tuning Language Models for Joint Rewriting and Completion of Code with Potential Bugs (2024)

Abstract: Handling drafty partial code remains a notable challenge in real-time code suggestion applications. Previous work has demonstrated shortcomings of large language models of code (CodeLLMs) in completing partial code with potential bugs. In this study, we view partial code as implementation hints and fine-tune CodeLLMs to jointly rewrite and complete partial code into functional full programs. We explore two strategies: one-pass generation and multi-pass iterative refinement. We construct new training and testing datasets using semantic-altering code transformations and iterative self-generations.We conduct comprehensive experiments over three representative open-sourced CodeLLMs – InCoder, CodeGen, and StarCoder.Results show that CodeLLMs fine-tuned using our approach achieve superior pass rates compared to the previous baselines across existing and newly-created benchmarks, effectively handle both potentially buggy and clean code, and largely preserve the integrity of the original partial implementations. We further present findings on the properties of the potential bugs we tested and on the design choices of our methods.

## introduction large language models of code (codellms), i.e. large language models trained on vast repositories of code, have demonstrated remarkable progress in program synthesis and other coderelated tasks @xcite @xcite . such models are commonly trained on finished and reviewed code, e.g. github crawls of open-source projects. however, one common application of such code intelligence is suggesting implementations for in-progress code, which tends to be less refined and more error-prone. previous work @xcite framed this inference challenge as "buggy-code completion" -completing code with potential bugs, where a potential bug is a piece of undesirable code in a code prefix with respect to some reference completion (see section 2). they showed that codellms, albeit achieving a high pass rate over code generation benchmarks, fail to generate functional code in the presence of potential bugs. figure 1 illustrates such an example. the marked potential bug in the partial code makes producing a functional completion a more challenging task. in this case, the completion model (codegen-2b) attempted to accommodate the potential bug by deviating from the reference completion, but the resulting completion is still functionally incorrect (panel a). besides, we also show the output of the completion model (codegen-2b) using only the problem description (panel b). we can observe that the generated completion code segment accurately addresses the problem, but it deviates from the given partial code in terms of function names and certain variables, which may disalign with user expectations. typically, users prefer suggestions that closely match their existing code structure and naming conventions. in this paper, we propose to fine-tune codellms to jointly rewrite and complete partial code. as shown in figure 1 , the finetuned codegen-2b gives a correct and full program (panel c). from the resulting completion, we can observe that the program aligns well with the given partial code while fixing the potential bugs. we hypothesize that a user's partial code offers valuable information about their intent. it first helps clarify the problem, providing a richer specification than mere descriptions. it at the same time helps narrow down which solution idea, among a possibly vastly diverse solution space, the user wants to use to address the problem. thirdly, it provides an implementation sketch for the solution idea, presumably making the task of finishing a functional program easier. it thus becomes more surprising to see that completing with the presence of potential """ you are given a list of deposit and withdrawal operationson a bank account that starts with zero balance. your task is to detect if at any point the balance of account falls below zeros, and at that point function should return true. otherwise, it should return false. """ def below_zero(operations): @xmath0: balance -= op if @xmath1: we hypothesize that a user's partial code offers valuable information about their intent. it first helps clarify the problem, providing a richer specification than mere descriptions. it at the same time helps narrow down which solution idea, among a possibly vastly diverse solution space, the user wants to use to address the problem. thirdly, it provides an implementation sketch for the solution idea, presumably making the task of finishing a functional program easier. it thus becomes more surprising to see that completing with the presence of potential bugs gives an even worse pass rate than generating code from only specifications @xcite . in this paper, we focus on tuning existing codellms with the task of generating a full functional programs given a problem specification and a code prefix, in which potential bugs may present. the intuition is to treat code prefixes as implementation hints and allow deviation from it. we explore two strategies of joint rewriting and completion of partial code: one-pass generation where a model directly attempt to generate a functional full program from a code prefix, and iterative refinement where a model iteratively refines its previous outputs. regarding the construction of the required training datasets, we design an automatic way to injecting multiple types of semantic-altering code transformations to code prefixes and use the original reference code as generation targets. we evaluate our models in terms of pass rates against b-humaneval and b-fixeval benchmarks from @xcite , as well as two newly constructed evaluation datasets b-mbpp and b foot_0 -mbpp. extensive results using codegen-(2b, 16b) @xcite , incoder-(1b, 6b) @xcite , and starcoder (15.5b) @xcite suggest that the iterative inference strategy outperforms the onepass generation strategy, which in turn surpasses previous post-hoc baseline mitigations @xcite in terms of pass rates on both buggy and reference code prefixes, and perseverance of the given 2 figure 1 : buggy-code completion with three different paradigms: completion with both the problem description (h), completion with only the problem description (h), and the code prefix (s) and joint rewriting with both the problem description (h) and the code prefix (s). a is incorrect, b is correct but disaligns with the given partial prefix (the function name and the variable names are different), and c is correct and align with the given partial prefix. in this paper, we focus on tuning existing codellms with the task of generating a full functional programs given a problem specification and a code prefix, in which potential bugs may present (section 2). the intuition is to treat code prefixes as implementation hints and allow deviations from it. we explore two strategies of joint rewriting and completion of partial code (section 3): onepass generation where a model directly attempts to generate a functional full program from a code prefix, and multi-pass iterative refinement where a model iteratively refines its previous outputs. regarding the construction of the required training datasets, we design an automatic way to injecting multiple types of semantic-altering code transformations to code prefixes and use the original reference code as generation targets (section 4). we evaluate our models in terms of pass rates against b-humaneval and b-fixeval benchmarks from @xcite , as well as two newly constructed evaluation datasets b-mbpp and b 2 -mbpp (section 4.4). extensive results (section 5) using codegen-(2b, 16b) @xcite , incoder-(1b, 6b) @xcite , and starcoder (15.5b) @xcite suggest that the multi-pass iterative inference strategy outperforms the one-pass generation strategy, which in turn surpasses previous post-hoc baseline mitigations @xcite in terms of pass rates on both buggy and reference code prefixes and in terms of preserving the given partial code. as illustrated in figure 1 , our approach not only provides a correct solution that aligns with the given partial code but also rectifies the potential bug, thus resulting in a more desirable outcome. contributions we (i) first attempt the joint rewriting and completion of code with potential bugs using two inference strategies facilitated by fine-tuning, (ii) construct 2 training datasets and 2 task in this section, we look into the task of transforming a code prefix s into a functional full program t, conditioning on a task specification h. motivated by the scenario of real-time code suggestion, where a user's in-progress code can be less refined and more error-prone than the final revised code, the code prefix here may contain potential bugs. a potential bug is a span u in a code prefix s, such that, compared to the reference program t := ŝ :: ĉ that implements h, s and ŝ only differs in u and that program t := s :: ĉ fails h. :: denotes code concatenation. in other words, u is a bug in program t. we call ŝ a reference prefix, ĉ a reference completion, and s a (code) prefix with potential bugs, or a (potentially) buggy prefix. note that, although we refer to s as "buggy prefix" for brevity, s is not buggy per se without a reference completion, because "bugginess" is only defined over a full program. the task was previously studied as "buggy-code completion" by @xcite . in particular, since the code prefix s manifests a user's implementation intent, a desideratum is to minimize deviations from s. we quantify this by the similarity between generated solution t and the reference solution t (see section 5.2), because the prefix-completion boundary is not always well-defined in the generated program t after the joint rewriting and completion. ## method we describe two inference strategies for joint rewriting and completion of code, one-pass generation and multi-pass iterative refinement, as well as how we fine-tune a codellm with such strategies. ## construction of datasets we create labeled instances for joint rewriting and completion of partial code using truncation and semantic-altering code transformations. section 4.1 -4.3 describe the construction of d 1 . section 4.4 describes testing datasets. ## analysis are buggy code prefixes recoverable? a partial code snippet, in itself, is not inherently buggy, as the notion of bugginess is ill-defined without the context of a complete program. to quantify the extent to which a partial code can be considered truly "buggy", we measure the pass@100 scores for the code completion results from a substantial large codellm (codegen-16b). if none of 100 generated solutions is able to pass all the test cases, it indicates that the partial code is impossible, or at least very difficult, to reach a correct solution by merely appending a suffix. in other words, these potential bugs are considered as "not recoverable". modifications over the given partial code may be necessary to reach a correct solution. figure 5 shows that even codegen-16b, a substantially large language model for code with up to 16.1b parameters, has lower pass@100 for the buggy partial code than that for the clean partial code. this suggests that obtaining a fully functional program merely by appending a suffix to the given buggy partial code is extremely challenging. to some extent, it indicates that rewriting the partial code with potential bugs fixed before the completion operation is necessary as to achieve a functional program. impact of different types of potential bugs figure 6 shows the results for different bug types obtained through three distinct methods. notably, @xmath2: k-remove , c-remove , and b-remove . these particular bug types are attributed to code removal operations. the bug fixer, realit (richter and wehrheim, 2022), employs a localization and repair module, which first identifies the bug's location and then seeks a suitable replacement token for repair. consequently, for the bugs arising from the removal of code segments or entire lines, the localization module struggles to pinpoint the bug. in cases where the code removal operation occurs at the last line of the partial code, there is no bug to fix, which explains why @xmath3. similar challenges are observed for @xmath4, as it also needs to identify lines of code containing bugs before utilizing the infilling operation. therefore, both of these post-hoc baseline methods have limitations in addressing specific bug types. by contrast, our fine-tuned model generates solutions from scratch while adhering to the given specification. it treats the provided partial code as a soft constraint to guide the generation process. this approach offers greater flexibility and wider applicability across various bug types. ## conclusion in this paper, we study the task of generating a complete functional program given a specification and an in-progress code prefix with potential bugs. to align codellms with the task of jointly rewriting and completing such code prefixes, we construct a new training dataset by injecting various potential bugs. beyond the standard fine-tuning and generation workflow, we apply an multi-pass iterative refinement strategy to enhance code generation performance. our experiments demonstrate that our methods outperform previous baseline methods on both existing evaluation benchmarks and the new testing datasets we have constructed.