SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Lakshya A Agrawal
t-lakagraphal@microsoft.com
Aditya Kanade
kanadeaditya@microsoft.com
Navin Goyal
navingo@microsoft.com
Shuvendu K. Lahiri
Redmond, United States
shuvendu.lahiri@microsoft.com
Sriram K. Rajamani
sriram@microsoft.com
Language models of code (LMs) work well when the surrounding code provides sufficient context. This is not true when it becomes necessary to use types, functionality or APIs defined elsewhere in the repository or a linked library, especially those not seen during training. LMs suffer from limited awareness of such global context and end up hallucinating.
Integrated development environments (IDEs) assist developers in understanding repository context using static analysis. We extend this assistance, enjoyed by developers, to LMs. We propose monitor-guided decoding (MGD) where a monitor uses static analysis to guide the decoding. We construct a repository-level dataset PRAGMATICCODE for method-completion in Java and evaluate MGD on it. On models of varying parameter scale, by monitoring for type-consistent object dereferences, MGD consistently improves compilation rates and agreement with ground truth. Further, LMs with fewer parameters, when augmented with MGD, can outperform larger LMs. With MGD, SantaCoder-1.1B achieves better compilation rate and next-identifier match than the much larger text-davinci-003 model.
We also conduct a generalizability study to evaluate the ability of MGD to generalize to multiple programming languages (Java, C# and Rust), coding scenarios (e.g., correct number of arguments to method calls), and to enforce richer semantic constraints (e.g., stateful API protocols). Our data and implementation are available at https://github.com/microsoft/monitors4codegen.
Language models of code (LMs), such as in Chen et al. (2021); Nijkamp et al. (2023); Allal et al. (2023) and many others, are revolutionizing code generation. Many commercial offerings based on LMs, like GitHub Copilot, Amazon Code Whisperer and Replit, are now available. The LMs work well when the surrounding code in the vicinity of generation provides sufficient context. This
Figure 1: Motivating example to illustrate monitor-guided decoding (MGD).
condition does not hold when it becomes necessary to use types, functionality or APIs defined in another module in the repository or an external library, especially those not seen during training. In the absence of awareness of such global context, the LMs end up hallucinating, e.g., using types defined in other files incorrectly. Further, devoid of the repository context, the LMs may lack awareness of other semantic information like the number of arguments required by a called method, or the ordering constraints on method calls (API protocols) to be followed. Often, such type and semantic information can come from artifacts generated at build-time, like Project Lombok (Iom, 2009) and ProtoBuf (pro, 2008), and therefore may not even be present as code context in the repository.
As an example, consider the partial code (method to be completed) in Figure 1(a). To complete this code, an LM has to generate identifiers consistent with the type of the object returned by ServerNode.Builder.newServerNode(). The method newServerNode and its return type, class ServerNode.Builder, are defined in another file. If an LM does not have information about the ServerNode.Builder type, it ends up hallucinating.
We show a completion generated by the OpenAI text-davinci-003 (Ouyang et al., 2022) and SantaCoder (Allal et al., 2023) models in the box decorated with  $\times$  in Figure 1(a). The completion uses identifiers host and port, which do not exist in the type ServerNode.Builder. The generated code therefore results in "symbol not found" compilation errors. The lack of awareness of other semantic information from the global context may result in other kinds of errors like compile-time errors (e.g., "actual and formal argument lists differ in length" on using wrong number of arguments) or runtime errors (e.g., IllegalStateExceptionException on violation of API protocols).
Integrated development environments (IDEs) have been at the forefront of assisting developers. Our inspiration is the use of static analysis by IDEs to bring the global context at the fingertips of developers. Many analyses are integrated in IDEs (Fuhrer, 2013) to infer and enforce semantic constraints on the code under development, e.g., resolving def-use, symbol references, and type hierarchies. Recently, there has been a rise in the use of Language Server Protocol (LSP) (lsp), which is an open industry standard of communication between IDEs and programming language specific tools like static analyzers and compilers, called Language Servers. There are a large number of Language Servers available, targeting most programming languages (lan, 2023), and providing a variety of syntactic and semantic information. In this work, we focus on the type-directed code completion analysis available through LSP in a language-agnostic manner, to provide guidance to an LM.
We propose a notion of monitors as a stateful interface between LMs and static analysis. A monitor observes the code generated by an LM and queries static analysis at pre-defined trigger points. The suggestions returned by the static analysis are converted to masks which are used for reshaping the logits (or equivalently, token-generation probabilities) produced by the LM in the subsequent decoding steps. We call our method monitor-guided decoding (MGD). Unlike an LM, a static analysis operates on the entire repository and its dependencies. While the LM generates completions by conditioning on the local context, the static analysis ensures consistency with the rest of the code in the repository. Through MGD, we bring the two together without the need to retrain the LM, and making a minor and modular addition to the decoding stage of the LM.
Figure 1(a) also shows the code generated by the SantaCoder model with MGD in the box decorated with  $\checkmark$ . This code makes use of identifiers that are actually defined in the class
ServerNode.Builder. It compiles and matches the ground truth. In comparison, the same SantaCoder model without MGD generates the erroneous code shown in the box decorated with  $\times$ .
Some recent approaches use static analysis (Shrivastava et al., 2022; Ding et al., 2022; Pei et al., 2023) or retrieval (Zhang et al., 2023) to extract relevant code fragments from the global context. These approaches expand the prompt (Shrivastava et al., 2022; Pei et al., 2023; Zhang et al., 2023) or require architecture modifications (Ding et al., 2022) and additional training (Ding et al., 2022; Pei et al., 2023). In comparison, we provide token-level guidance to a frozen LM by invoking static analysis on demand. Our method is complementary to these approaches as they condition the generation by modifying the input to the LM, whereas we apply output-side constraints by reshaping the logits.
We make the following contributions in this paper:
- Monitor-Guided Decoding (MGD) as a stateful interface between LMs and static analysis. A monitor watches the LM generating code, queries static analysis in the background, and uses the information from the static analysis to effectively guide the decoding stage of LMs.
- PRAGMATICCODE, a publicly-released dataset of Java code repositories complete with their development environments and dependencies.
- Instantiation of MGD for generating code having type-consistent identifier dereferences.
- Large scale evaluation on PRAGMATICCODE showing: (1) MGD consistently improves the ability of an LM (across varying parameter scales) to generate type-consistent identifiers, compilation rates and agreement with ground truth. (2) Further, LMs with fewer parameters, in conjunction with MGD, can outperform larger LMs. For instance, SantaCoder-1.1B with MGD achieves better compilation rate and next-identifier match than the much larger text-davinci-003 model, when both have a budget of 1 generation each. With a budget of 4 generations, it also surpasses agreement with ground truth of text-davinci-003. (3) We also evaluate how MGD complements different prompt augmentation and decoding strategies.
- Microbenchmark to demonstrate generalizability of MGD to different (1) programming languages, (2) coding scenarios, and (3) use of other static analysis techniques for guiding with rich semantic properties. Notably, we demonstrate that small-LMs can be guided to adhere to rich static properties like satisfaction of stateful API protocols.
- We open source our implementation and provide an extensible Python library called multilspy with static analysis bindings for multiple languages over LSP, suitable for monitoring of LMs. It can be used for other AI4Code scenarios as well.
2 Monitor-Guided Decoding
Background. Static analysis of code (Nielson et al., 2015) is used widely in industry in various applications such as in detecting bugs and optimizing code. While analysis is usually performed on complete code, IDEs have long applied static analysis on incomplete code under development (Reps et al., 1983; Dagenais & Hendren, 2008), using algorithms for incremental parsing and semantic analysis (e.g., type inference) of partial code (Hedin, 1992; Wagner, 1997; Maddox III, 1997). These analyses have now become a standard part of language servers.
A key abstraction used in analysis of partial code is partial abstract syntax trees (ASTs) with special nodes to indicate incomplete parts of code. These ASTs are further decorated by semantic information through attribute grammars (Reps et al., 1983) that decorate each AST node with attributes that capture the static semantics not captured in the context-free grammar (such as consistency of types among expressions in an assignment). This can range from computing the type-hierarchy for object oriented languages, binding the variables in AST nodes to their declarations, resolving the types of expressions as well as computing the def-use relationships for resolved AST nodes (Fuhrer, 2013).
Figure 1(b) shows the partial AST for the incomplete code in Figure 1(a). All the statements upto the incomplete return statement are completely parsed and subtrees corresponding to them are constructed. The subtree for the return statement includes a node [UNKNOWN] indicating the incomplete part. As shown in Figure 1(b), an incremental semantic analysis resolves the type of the partial expression ServerNode.Builder.newServerNode() to ServerNode.Builder. Later, we show how to use this type information to construct a monitor, which can then be used to guide an LM to generate type-consistent identifier completions.
Basic Concepts and Notation. Consider an LM  $L_{\theta}$  operating on a vocabulary  $V$ . Let  $x_{1}, \ldots, x_{n}$  be the partial code that has been generated by the LM and  $x_{n+1}$  be a candidate next token. Though a vanilla (auto-regressive) prompt would consist only of  $x_{1}, \ldots, x_{n}$ , today many approaches augment it with additional information. We use  $p$  to indicate this additional prompt, e.g., the suffix information used in fill-in-the-middle prompting (Donahue et al., 2020; Fried et al., 2022; Bavarian et al., 2022).
A property  $\varphi$  specifies the constraints that a piece of code needs to satisfy. A static analysis  $A_{\varphi}$  checks whether the partial code satisfies  $\varphi$ . If yes, it returns suggestions to extend it so that the extended code continues to satisfy  $\varphi$ . The static analysis works on the repository context  $C$  which not only includes code spread across multiple files in the repository, but also external dependencies and intermediate artifacts (e.g., code bindings) generated during the build process. Such repository context is often very large, diverse and complex. Directly including it as an input to the LM will result in bloating and pass the burden of distilling useful information from it to the LM.
A monitor  $M_{\varphi}$  for a property  $\varphi$  is a tuple  $(A_{\varphi}, s_0, S, \text{pre}, \text{update}, \text{maskgen})$ . The monitor starts in the wait state  $s_0$ . If the partial code satisfies the pre-condition pre then the monitor is triggered and it invokes  $A_{\varphi}$  on the partial code. The monitor maintains the suggestions returned by  $A_{\varphi}$  in its state and uses them to guide sampling of the next token  $x_{n+1}$ . Let  $S$  be the set of states and  $s, s' \in S$  respectively be the current and next states of the monitor. With each sampled token  $x_{n+1}$ , the monitor updates its state using a function update  $(s, x_{n+1})$  to a new state  $s'$ , which tracks the residual suggestions after the token  $x_{n+1}$  is output. When the suggestions are exhausted, it reverts to the wait state  $s_0$ . We explain the function maskgen below.
Decoding Process. Usually, the next token  $x_{n+1}$  can be any token from the vocabulary  $V$ , sampled based on the logits  $\ell$  determined by the LM. Unlike the usual decoding, in monitor-guided decoding, we supervise the code generation using a monitor  $M_{\varphi}$  for a property  $\varphi$ . We denote the composition of  $L_{\theta}$  and  $M_{\varphi}$  by  $L_{\theta}||M_{\varphi}$ , meaning that both the LM and the monitor are running concurrently and sampling the tokens jointly. The decoding is conditioned on the partial code  $x_1, \ldots, x_n$ , the repository context  $C$ , the prompt  $p$  and the current state  $s$  of the monitor.
Eq. (1) states that whenever the monitor is in the wait state  $s_0$ , we sample  $x_{n + 1}$  as per the logits  $\ell$  determined by the LM (Eq. (2)). Otherwise, the logits are combined with a mask  $m$  using a function  $\oplus$  such that if  $m[x] = 0$  then  $\ell [x]$  is reset to a large negative value  $-K$  and is left unchanged otherwise. This mask is computed by the function maskgen in Eq. (3) guided by the current state  $s$  of the monitor. Eq. (4) defines how the state of the monitor evolves. When the pre-condition  $\mathsf{pre}(s;x_1,\ldots ,x_n)$  evaluates to true, the next state  $s^{\prime}$  of the monitor is determined by the suggestions returned by the static analysis  $A_{\varphi}$ . Otherwise, it is determined by the update function.
$$
\left(L _ {\theta} \mid \mid M _ {\varphi}\right) \left(x _ {n + 1} \mid x _ {1}, \dots , x _ {n}; C, p, s\right) = \left\{ \begin{array}{l l} \operatorname {s o f t m a x} (\ell) \left[ x _ {n + 1} \right] & \text {i f} s = s _ {0} \text {i s t h e w a i t s t a t e} \\ \operatorname {s o f t m a x} (\ell \oplus m) \left[ x _ {n + 1} \right] & \text {o t h e r w i s e} \end{array} \right. \tag {1}
$$
$$
\ell = L _ {\theta} \left(\cdot \mid x _ {1}, \dots , x _ {n}; p\right) \tag {2}
$$
$$
m = \operatorname {m a s k g e n} (s, V) \tag {3}
$$
$$
s ^ {\prime} = \left\{ \begin{array}{l l} A _ {\varphi} \left(x _ {1}, \dots , x _ {n}; C\right) & \text {i f} s = s _ {0} \wedge \operatorname {p r e} \left(s; x _ {1}, \dots , x _ {n}\right) \\ \operatorname {u p d a t e} \left(s, x _ {n + 1}\right) & \text {o t h e r w i s e} \end{array} \right. \tag {4}
$$
The specifics of the monitor state, and the pre, update and maskgen functions depend on the static analysis  $A_{\varphi}$  used by the monitor. Our formulation is general and even allows combining multiple static analyses by taking a product of the state-spaces of their respective monitors. In the following, we discuss a specific instantiation of this framework of monitor-guided decoding.
Monitoring for Use of Type-Consistent Identifiers. When an object obj of a type  $T$  is dereferenced, the next token (or more generally, the sequence of subtokens) should refer to an identifier of a field or method defined by the type  $T$ . It can otherwise result in a "symbol not found" error. The type  $T$  could be defined in another package, imported file or in a library. Unless  $T$  comes from a popular library seen during training, the LM may not have knowledge about  $T$ .
Our monitor  $M_{\varphi}$  is triggered when the partial code  $x_{1}, \ldots, x_{n}$  ends with a partial object dereference expression "obj." where ". ." is the dereference operation. This is the pre-condition pre we use. We employ a static analysis  $A_{\varphi}$  which returns all the type-consistent identifiers that can be referenced through obj. For this,  $A_{\varphi}$  performs a global analysis over the partial code, imported files, libraries used, and class hierarchies to infer the type  $T$  of obj and the identifiers accessible through  $T$ . The set of type-consistent identifiers returned by  $A_{\varphi}$  forms the state of the monitor (see Eq. (4)).
Given a state  $s$  and the vocabulary  $V$  of the LM, maskgen identifies all the tokens in  $V$  that are consistent with the suggestions in  $s$ . The identifiers returned by the static analysis are tokens as per the programming language, whereas the vocabulary  $V$  may use its own space of subtokens (Schuster & Nakajima, 2012; Kudo & Richardson, 2018). The mask  $m$  (Eq. (3)) is generated by string matching. For all tokens  $t \in V$  that form prefixes of the identifiers in  $s$ , the mask value  $m[t]$  is set to 1, indicating that they can be sampled. Let  $E$  be the set of special symbols that indicate end of an identifier name, e.g., the symbol '(' in 'getName(' or ', in 'x','. Let  $w$  be a string in  $s$  and  $\Sigma$  be the set of all possible characters. If a token  $t \in V$  matches the regular expression  $w \cdot E \cdot \Sigma^*$  then its mask value  $m[t]$  is also set to 1. For all other tokens  $t$  in  $V$ , the mask value  $m[t]$  is set to 0.
Let  $x_{n+1}$  be the next token sampled according to the second equation in Eq. (1). If  $x_{n+1}$  contains a symbol from  $E$ , indicating that a complete identifier name has been generated in accordance with the set returned by  $A_{\varphi}$ , the monitor reverts to the wait state to wait for the next trigger. Otherwise, the token  $x_{n+1}$  must be a prefix of a member of  $s$ . The update function removes the members in  $s$  that are not prefixed by  $x_{n+1}$ , and those prefixed by  $x_{n+1}$  are updated by pruning the prefix string  $x_{n+1}$ . The resulting set of character strings forms the next state  $s'$  (see the second equation in Eq. (4)). If  $A_{\varphi}$  returns an empty set to start with, we abandon the current run. Note that a single identifier may need to be generated using multiple tokens. Figure 4 (see Appendix A) shows the interaction between the LM and the monitor for the example in Figure 1, and specifically illustrates how the complete identifier names suggested by the static analysis are gradually pruned by the monitor to corresponding suffixes in each successive state as the prefixes get generated as tokens.
Dataset Creation. In order to evaluate MGD, we need real-world repositories with their build environments and dependencies. Most published datasets are standalone, with the exception of CoderEval (Yu et al., 2023) and PyEnvs (Pei et al., 2023), both of which are not publicly available at the time of this writing. Hence we curated PRAGMATICCODE, a dataset of real-world open-source Java projects complete with their development environments and dependencies. We ensure that these repositories were released publicly only after the determined training dataset cutoff date (31 March 2022) of the models which we use to evaluate MGD.
From PRAGMATICCODE, we identify a set of method-level completion task instances, creating DOTPROMPTs as a method-level code completion benchmark. Each testcase in DOTPROMPTs consists of a prompt upto a dereference location (using the “.” operator in Java) within a target method, and the task is to complete the remainder of the method. We ensure sufficient complexity in the identified target methods in DOTPROMPTs by including methods that satisfy a set of complexity filters (e.g., the method should have at least 7 lines of code) described in detail in appendix B. Overall, PRAGMATICCODE consists of 100 repositories, and DOTPROMPTs consists of 1420 methods and 10538 dereference prompts. Appendix B gives further details.
Models. We study the effect of performing MGD on code generation with the HuggingFace Transformers (Wolf et al., 2020) implementation of Salesforce CodeGen family of models (CodeGen $\{350M, 2B, 6B\}$ -Multi, abbreviated as CG- $\{350M, 2B, 6B\}$  hereafter) (Nijkamp et al., 2023) and BigCode SantaCoder-1.1B (SC or SantaCoder hereafter) (Allal et al., 2023). We also evaluate OpenAI text-davinci-003 (TD-3 hereafter) with and without MGD, available on Azure.
Prompting Strategies. We study the effect of different prompt augmentation techniques when combined with MGD: (1) Standard: Include the local file content up to the dereference point and truncate from left to fit the prompt budget. (2) classExprTypes: For a given target method belonging to a class C, identify the type of all expressions occurring in C (after masking out the target method to prevent leakage) and include the concatenated file contents for the type definitions of all the identified files, truncating from the left as necessary. We assign a budget of  $20\%$  tokens of total prompt budget to classExprTypes. (3) RLPG: Use the prompt augmentation technique proposed in Shrivastava et al. (2022). We use their released source code and model checkpoints to adapt RLPG to DOTPROMPTS.
Decoding Strategies. We experiment with two decoding strategies: (1) Autoregressive: for left-to-right decoding. (2) Fill-in-the-middle: Use fill-in-the middle (FIM) setting implemented in SantaCoder (Allal et al., 2023).
Table 1: Summary of results with a budget of 6 generations per model: The numbers in parentheses are relative improvements of the "-MGD" configuration over the respective base model.
Config\(\backslash\) (Metric, score@6)	CR	NIM	ISM	PM
CG-350M	52.43	76.94	31.86	26.86
CG-350M-MGD	65.37 (24.69%)	83.80 (8.92%)	34.31 (7.70%)	28.69 (6.82%)
CG-2B	57.01	81.11	36.38	30.95
CG-2B-MGD	70.91 (24.38%)	87.32 (7.66%)	39.03 (7.29%)	33.06 (6.82%)
CG-6B	58.64	81.55	37.17	31.69
CG-6B-MGD	72.28 (23.25%)	87.35 (7.11%)	39.55 (6.41%)	33.56 (5.89%)
SC	59.97	82.40	38.14	32.10
SC-MGD	73.03 (21.77%)	88.42 (7.31%)	40.69 (6.68%)	34.25 (6.72%)
SC-classExprTypes	64.57	84.91	39.67	33.55
SC-classExprTypes-MGD	75.01 (16.18%)	89.37 (5.25%)	41.56 (4.78%)	34.98 (4.26%)
SC-RLPG	66.39	85.42	42.35	36.21
SC-RLPG-MGD	78.14 (17.70%)	89.89 (5.23%)	44.47 (5.00%)	37.97 (4.87%)
SC-FIM	68.23	85.56	42.22	36.12
SC-FIM-MGD	80.19 (17.52%)	89.89 (5.07%)	44.50 (5.39%)	37.91 (4.95%)
SC-FIM-classExprTypes	70.97	86.99	42.67	36.36
SC-FIM-classExprTypes-MGD	80.33 (13.18%)	90.42 (3.94%)	44.18 (3.54%)	37.75 (3.82%)
TD-3	62.66	86.18	44.97	38.77
TD-3-MGD	74.26 (18.52%)	91.19 (5.81%)	47.33 (5.24%)	39.94 (3.03%)
Metrics. We use the following metrics to measure the quality of generated code: (1) Compilation Rate (CR): We replace the ground truth method with the generated method in the context of the complete repository and invoke a clean build. We assign a score of 1 if the compilation succeeds, and 0 otherwise. (2) Match with the ground truth: We use three specific metrics to measure how closely the generation matches ground truth, namely (a): Next Identifier Match (NIM): If the first Java token generated by the LM matches with the ground truth, we assign a score of 1, 0 otherwise; (b) Identifier Sequence Match (ISM): Longest prefix match between the ordered set of identifier names in the ground truth and generated completion, normalized by the number of identifiers in the ground truth; and (c) Prefix Match (PM): Longest prefix match between the ordered set of Java tokens (as obtained from a Java Lexer) between the ground truth and generated completion, normalized by the number of tokens in the ground truth. Except NIM, all other metrics - namely CR, ISM and PM - evaluate the complete method-level generation by the model.