Source-linked AI summary
Repository-Level Prompt Generation for Large Language Models of Code
Disha Shrivastava, Hugo Larochelle, Daniel Tarlow
TL;DR
Prompt design for code LLMs remains poorly understood, especially for incorporating repository-specific context. RLPG learns example-specific prompts from repository-level proposals and achieves up to 36% relative improvement over Codex with an oracle.
Problem
Effective task-specific prompt design for code LLMs remains poorly understood, motivating study of useful repository-level contexts and automatic prompt selection.
Method
RLPG learns example-conditioned prompts by selecting proposal contexts from repository structure and relevant files, requiring only black-box access to the LLM.
Results
Up to 36% relative improvement over Codex was achieved on single-line code autocompletion by an oracle using the proposed prompt proposals.
Takeaways & Limitations
RLPG provides a way to adapt prompts to repository-specific context without repeatedly fine-tuning the language model.
Takeaways & Limitations
Because repositories are often long, including an entire repository in the prompt is impractical; 70.22% of deduplicated repositories exceeded 32k tokens.
Abstract
from arXiv · showhide
With the success of large language models (LLMs) of code and their use as code assistants (e.g. Codex used in GitHub Copilot), techniques for introducing domain-specific knowledge in the prompt design process become important. In this work, we propose a framework called Repo-Level Prompt Generator that learns to generate example-specific prompts using prompt proposals. The prompt proposals take context from the entire repository, thereby incorporating both the structure of the repository and the context from other relevant files (e.g. imports, parent class files). Our technique doesn't require any access to the weights of the LLM, making it applicable in cases where we only have black-box access to the LLM. We conduct experiments on the task of single-line code-autocompletion using code repositories taken from Google Code archives. We demonstrate that an oracle constructed from our prompt proposals gives a remarkably high relative improvement of 36% over Codex, showing the quality of these proposals. Further, we show that when we train a model to predict a prompt proposal, we can achieve significant performance gains over Codex and other baselines. We release our code, data, and trained checkpoints at: \url{https://github.com/shrivastavadisha/repo_level_prompt_generation}.
1. Introduction
The paper proposes RLPG, which generates example-specific prompts using repository structure and relevant cross-file context. It targets single-line code autocompletion while requiring only black-box access to the language model.
- Motivation and contribution: RLPG learns to generate prompts conditioned on each example through a set of repository-level prompt proposals.Prompt proposals specify where to retrieve context from the repository and what information to include.
- Motivation and contribution: The framework incorporates repository structure and relevant context from files such as imports and parent classes into generated prompts.This enables prompt engineers to inject domain expertise into prompt design for repository-specific software.
- Motivation and contribution: RLPG assumes only black-box access to the LLM and avoids repeatedly fine-tuning model weights when adapting to proprietary or niche repositories.This addresses settings where model weights and training data are unavailable while still adapting prompts to repository-level context.
- Evaluation setting: The evaluation focuses on single-line code autocompletion, predicting the blanked-out portion from an imagined cursor to the end of the line.A prompt proposal classifier selects repository context, and a prompt composer combines it with the current file’s default Codex context.
- Evaluation setting: An oracle constructed from RLPG’s prompt proposals yields a substantial improvement over Codex on single-line code autocompletion.The introduction reports this oracle result as evidence of the quality of the generated prompt proposals.
2. Repo-Level Prompt Generator (RLPG)
RLPG generates repository-aware prompt proposals by replacing parts of Codex’s default context with selected repository context, then predicts useful proposals and composes them with Codex context. Its proposals combine ten sources with seven context types, while classification treats successful proposals as multi-label predictions and supports two classifier variants.
- 2.1. Repo-Level Prompt Proposals: RLPG replaces part of Codex’s default context with repository context selected through prompt proposals designed to capture common coding patterns.A proposal specifies both a prompt source and a prompt context type.
- 2.1. Repo-Level Prompt Proposals: Ten prompt sources determine where proposal code comes from, including the current file and the parent-class file.The current file can contribute code after the hole, while the parent-class file captures inherited method context.
- 2.1. Repo-Level Prompt Proposals: 63 prompt proposals result from combining prompt sources with prompt context types, although applicability depends on the target hole.Some sources can provide a single file or ranked files until the proposal’s maximum context length is exhausted.
- 2.1. Repo-Level Prompt Proposals: Seven prompt context types determine what code to extract, including post lines, identifiers, type identifiers, field declarations, string literals, and method names with signatures.The supplied passage lists these context types as proposal components.
- 2.2. Prompt Proposal Classifier (PPC): The prompt proposal classifier predicts proposals that produce an exact match between the predicted hole and target hole, using multi-label binary classification.The default Codex context is included as one of the prompt proposals because multiple proposals may succeed for one hole.
- 2.2. Prompt Proposal Classifier (PPC): Training labels are multi-hot vectors formed by testing each applicable proposal with Codex, marking proposals successful when their predicted hole exactly matches the target.Masking and binary cross-entropy consider only proposals applicable to each hole.
- 2.2. Prompt Proposal Classifier (PPC): RLPG-H represents a hole window with a pretrained model, using the [CLS] hidden state as its context representation.The window includes two lines before the hole, code up to the hole, and two lines after it.
- 2.3. Prompt Composer: The prompt composer combines selected proposal context with Codex context under a fixed prompt length using dynamic allocation.Unused proposal-context capacity is assigned to the default Codex context when proposal context is shorter than its allocation.
3. Experiments and Results
Experiments evaluate repository-level prompt proposals and learned RLPG selection against Codex and retrieval baselines. Prompt proposals improve completion performance, with gains extending across exact-match, edit-distance, and cross-model evaluations.
- Computational Complexity and Scalability of RLPG: Approximately 230k Codex queries were used to obtain training and validation labels, while RLPG-R training required 9.19 minutes per epoch on one Tesla V100 GPU.RLPG-R has 3.6M parameters, compared with Codex’s 175B parameters, and uses 1.63 queries per target hole for label collection.
- RQ1 - Performance of Prompt Proposals: 36% relative improvement over Codex was achieved by the prompt-proposal oracle on the test split.The oracle combines successful prompt-proposal contexts with Codex’s default context and represents an upper bound because its information is unavailable at inference.
- RQ2 - Performance of PPC: All RLPG variants and the fixed prompt proposal significantly improve performance over Codex, while random baselines are worse or on par with Codex.RLPG-based methods outperform the fixed prompt proposal, and combining BM25 retrieval with prompt-proposal contexts improves performance substantially.
- Performance based on Prompt Proposals: The current file is the most successful prompt source, followed closely by sibling files and similar-name files.Every prompt source has a non-zero chance of success when applicable.
- Edit Distance as a Metric: 26.62% relative improvement was achieved by RLPG-H over Codex on normalized character-level edit distance.Lower normalized character-level edit distance is better, and both RLPG variants show significant relative improvements over Codex.
- Cross-model Evaluation: 10.87% and 10.95% relative improvements were achieved by RLPG-H and RLPG-R, respectively, over prior in-file context when evaluated with code-cushman-001.The RLPG models were trained using labels from code-davinci-001, then their predicted contexts were used to prompt code-cushman-001.
4. Related Work
Prior work spans multiple architectures for code LLMs, repository-level context modeling, and automatic prompt generation through learned continuous prompts.
- LLMs for Code: Code LLMs include decoder-only, encoder-only, and encoder-decoder architectures with distinct generation or encoding objectives.Decoder-only models generate code left to right, while encoder-only models use masked language modeling.
- Repo-Level Info: Repository-level methods exploit directory locality, parent classes, relational graphs, and API dependencies to incorporate relevant project context.These approaches support tasks including comment generation, context pruning, code summarization, and code generation.
- Prompt Generation: Automatic prompt generation in NLP includes continuous or soft prompts represented in a language model’s latent space.Prefix-Tuning learns a prefix added to the language model.
5. Discussion
The discussion presents RLPG as a black-box framework for generating repository-informed prompts while cautioning that LLM-based code-completion systems can produce insecure code or expose sensitive information.
- Limitations and deployment: Code-completion systems used with LLMs should be deployed cautiously because generated code may be insecure or contain sensitive information.The discussion also notes the introduction of larger-context LLMs, including GPT-4 with support for 32k tokens.
- Contributions: RLPG automatically generates example-conditioned prompts without requiring access to the LLM’s weights.It uses repository structure and context from other files through easy-to-understand prompt proposals.
- Future work: RLPG uses repository structure and other-file context through prompt proposals, taking context from only one proposal in this work.The authors identify composing prompts from multiple proposals as future work.
A. Dataset Creation Details … B. Prompt Proposal Details
The dataset was built from Java repositories and enriched with cached repository-level metadata for generating prompt proposals. Proposals select and format relevant code context using file-ranking heuristics, context types, truncation strategies, and tested variations.
- A.1. Creation of Hole Completion Data: Java hole-completion data came from Google Code repositories, deduplicated against GitHub by repository name, with unmatched archives downloaded and source structure preserved.Repository licenses were then identified through LICENSE files or license-related keywords.
- A.2. Creation of Data for Repo-Level Prompt Proposals: Repository-level proposal data used tree-sitter for Java 9 and stored class, import, sibling, and similar-name file relationships for each repository file.Import-use positions were also recorded to support heuristic file ranking.
- A.2. Creation of Data for Repo-Level Prompt Proposals: Metadata was computed once per repository and cached, after which the prompt source selected a file or ranked file list from parse-tree and repository metadata.The selected files supplied prompt context types including MN, MNB, SL, I, TI, and FD.
- B.1. Ranking of files based on prompt source: Prompt-source ranking selected current or parent-class files singly, while import and sibling files were ranked using import-use proximity and shared-import information.Table 6 specifies these file-selection rules for each prompt source.
- B.2. Examples of Prompt Context Type: Prompt context types included post lines, identifiers, type identifiers, method names, and method names with bodies, illustrated through code examples from the paper’s figure.Post lines extend from the specified line to the file end, while identifier context captures variable names and type identifiers capture identifier types.
- B.3. Truncation Strategies for Prompt Proposal Context: When proposal context exceeded its allocation, truncation used front or back schemes selected by prompt source and context type.For Current, non-PL context began after the hole and fell back to before it only when necessary, because Codex already supplied preceding code.
- B.4. List of Prompt Proposals: The proposed repo-level prompt proposals are enumerated in Table 7.The supplied passage identifies the table as the list of proposed prompt proposals but does not reproduce its entries.
- B.5. Other Prompt Proposal Variations: Additional variations changed class-name formatting, context separators, selected context types, top-k ordering, and whether proposal context was commented, labeled with methods, or passed unchanged.Separators included spaces and newlines; class_name formatting applied to MN, I, TI, FD, and SL, while class_method_name applied to MNB and comment formatting was not useful.
C. Implementation Details … C.3. Baselines
The implementation uses CodeBERT-based representations and specified neural architectures for RLPG-H and RLPG-R, while the baselines include random repository-context selection and nearest-neighbor similarity.
- C.1. RLPG-H: RLPG-H uses Adam with a 3e-4 learning rate, batch size 64, and CodeBERT representations of dimension 768 for the hole window.Its parameter matrices are W 1 ∈R512×768, b1 = 512, W 2 ∈R63×512, and b2 = 63.
- C.2. RLPG-R: RLPG-R uses Adam with a 3e-4 learning rate, batch size 64, and CodeBERT representations of dimension 768 for the hole window and prompt proposal context.The model includes multiheaded attention.
- C.2. RLPG-R: The multiheaded-attention formulation identifies dk as the key dimension and uses W Q projection matrices, W O, and τ attention heads.The passage describes W O as the linear projection combining the heads.
- C.2. RLPG-R: RLPG-R uses four attention heads with dk = dq = dv = 32 and dmodel = 768.It also specifies Wr ∈R63×768 and bp = 63.
- C.2. RLPG-R: The RLPG-R G module combines dropout, a residual connection, layer normalization, and dense layers with dimensions 2048 × 768 and 768 × 2048.The sequence also includes a ReLU activation and another dropout layer.
- C.3. Baselines: The random baseline selects a repository file and line randomly, then uses the lines from that point to the file’s end as context.For the current file, the hole window is excluded; nearest-neighbor similarity uses the dot product between CodeBERT representations of the hole window and context.
D. Additional Results · D.1. Hole-wise and Repo-wise results
The section compares methods using success rate averaged across completion holes and repositories. Repo-wise evaluation is independent of repository size, while Table 8 reports both aggregation schemes on test data.
- D.1. Hole-wise and Repo-wise results: Table 3 reports method performance averaged across all completion holes.
- D.1. Hole-wise and Repo-wise results: The hole-wise metric aggregates results over individual holes.
- D.1. Hole-wise and Repo-wise results: Table 3 also reports performance averaged across individual repositories.
- D.1. Hole-wise and Repo-wise results: The repo-wise metric is independent of repository size.
- D.1. Hole-wise and Repo-wise results: Table 8 presents hole-wise success rate for different methods on the test data.
- D.1. Hole-wise and Repo-wise results: Table 8 presents repo-wise success rate for different methods on the test data.
D.2. Ablation on Performance based on Prompt Proposal
The ablation finds that post lines is the most useful prompt context type by mean success rate when applicable, although other context types contribute significantly; after normalization across proposals, it is no longer dominant.
- Performance by prompt context type: Post lines has the highest mean success rate among applicable prompt context types.Figure 3 evaluates success only when each prompt context type applies.
- Performance by prompt context type: Other prompt context types contribute less than post lines but remain significant, highlighting the importance of each type.
- Normalized performance: Normalized success rates compare prompt proposal sources and context types relative to performance across prompt proposals.Figure 4 reports prompt sources on the left and prompt context types on the right.
- Normalized performance: Post lines is no longer the most dominant context type when relative performance is considered.
D.3. Performance on non-immediate Post Lines … D.7. Performance on individual repositories
Across these evaluations, repository-level prompt proposals remain useful under harder completion settings, varying composition and context lengths, while truncation and repository-level differences affect performance. The experiments also examine proposal composition, truncation rates, and success rates across individual repository splits.
- D.3. Performance on non-immediate Post Lines: The non-immediate post-line experiment tests whether post-line context remains useful when completion extends beyond the current line.It skips three lines after the target hole before providing post-line context.
- D.3. Performance on non-immediate Post Lines: Starting four lines after the target hole causes a slight performance deterioration, but the drop is not significant.This setting approximates predicting the blanked portion plus the next three lines for harder multi-line code autocompletion.
- D.4. Composition of prompt proposals: Composition uses the top-l prompt proposals ranked by decreasing RLPG probability and allocates total context length according to normalized probabilities.The experiment evaluates two RLPG versions composing prompt context from l proposals.
- D.5. Effect of Context Length: With half of Codex’s available context length, the oracle and fixed prompt proposal retain improvement over Codex, although relative gains are smaller than with larger context.The comparison highlights the value of prompt proposals while showing that context length influences relative gains.
- D.6. Effect of Truncation: Truncation rates are measured separately for default Codex context, prompt-proposal context, either context, and both contexts.The results suggest that code-completion models supporting longer context lengths can be useful.
- D.7. Performance on individual repositories: The individual-repository analysis reports success rates for different methods across training, validation, and test splits.Tables 13, 14, and 15 provide the repository-level results for these three splits.
- D.7. Performance on individual repositories: Repo-wise averages are computed by averaging repository-level column values, whereas hole-wise averages weight repository results by the total number of holes.The hole-wise calculation aggregates successful holes after multiplying each repository’s method result by its hole count.
E. Analysis of Sample Cases · E.1. Positive Cases · E.2. Negative Cases
The sample-case analysis shows RLPG correcting Codex when repository-derived prompt context exposes target-relevant information, while additional context can also confuse predictions. Positive examples include exact target fragments in proposal context; negative examples show incorrect predictions caused by misleading extra information.
- E. Analysis of Sample Cases: RLPG selected a proposal containing method names and bodies from the imported MaximizingGibbsSampler.java file.The imported-file context included InitializeToAssignment, while the default Codex context contained CurrentAssignme.
- E.1. Positive Cases: The analysis presents examples where RLPG produced the correct prediction while Codex failed.
- E.1. Positive Cases: When part of the target hole appeared exactly in proposal context, RLPG predicted fragments such as Propagation(int numVars) instead of Codex’s Propagation().
- E.1. Positive Cases: Other positive cases included RLPG predicting tersFromFile(String filename) {, als("dampingFactor")) {, and stem.exit(1); where Codex predicted different fragments.
- E.2. Negative Cases: The paper notes that extra information from prompt-proposal context can cause confusion and produce incorrect predictions.
- E.2. Negative Cases: In negative examples, RLPG predicted an hasConverged_; and _[i][j] = -Double.MAX_VALUE; while Codex predicted an converged_; and _[i][j] = 0;.