PaperlayerView exact revision on arXiv

Source-linked AI summary

RepoFusion: Training Code Models to Understand Your Repository

Disha Shrivastava, Denis Kocetkov, Harm de Vries, Dzmitry Bahdanau, Torsten Scholak

arXiv:2306.10998v1cs.LGcs.AIcs.PLcs.SE

TL;DR

Code models can struggle with repository context, particularly in repositories unseen during training. RepoFusion trains models to combine repository context for single-line completion, outperforming much larger next-token models and coming close to StarCoderBase.

  • Problem

    Code models can struggle to use repository-specific context, a challenge that is more pronounced for proprietary or work-in-progress repositories unseen during training.

    Source: §1 · pp. 1–2

  • Method

    RepoFusion adapts Fusion-in-Decoder to train code models to combine retrieved repository contexts with surrounding-file context for single-line completion.

    Source: §1 · p. 2

  • Results

    RepoFusion significantly outperformed larger next-token models such as CodeGen-16B and achieved a success rate 1.3% below StarCoderBase.

    Source: §1 · p. 2, §3.3 · p. 7

  • Takeaways & Limitations

    The results show repository-context training can let smaller code models surpass larger models without repository context in single-line completion.

    Source: §5 · p. 10

    Inference scales linearly with the number of repo contexts, leading to slower inference for larger context counts.

    Source: §5 · p. 9

Abstract

from arXiv · show

Despite the huge success of Large Language Models (LLMs) in coding assistants like GitHub Copilot, these models struggle to understand the context present in the repository (e.g., imports, parent classes, files with similar names, etc.), thereby producing inaccurate code completions. This effect is more pronounced when using these assistants for repositories that the model has not seen during training, such as proprietary software or work-in-progress code projects. Recent work has shown the promise of using context from the repository during inference. In this work, we extend this idea and propose RepoFusion, a framework to train models to incorporate relevant repository context. Experiments on single-line code completion show that our models trained with repository context significantly outperform much larger code models as CodeGen-16B-multi ($\sim73\times$ larger) and closely match the performance of the $\sim 70\times$ larger StarCoderBase model that was trained with the Fill-in-the-Middle objective. We find these results to be a novel and compelling demonstration of the gains that training with repository context can bring. We carry out extensive ablation studies to investigate the impact of design choices such as context type, number of contexts, context length, and initialization within our framework. Lastly, we release Stack-Repo, a dataset of 200 Java repositories with permissive licenses and near-deduplicated files that are augmented with three types of repository contexts. Additionally, we are making available the code and trained checkpoints for our work. Our released resources can be found at \url{https://huggingface.co/RepoFusion}.

1 Introduction

Code models can struggle with repository-specific context, especially in unseen or changing codebases. RepoFusion trains models to combine repository context for single-line code completion.

1 Introduction

  • RepoFusion learns to combine relevant repository cues to improve code predictions, addressing the difficulty of integrating multiple contexts without specialized training.Repository dependencies such as imports and parent classes can inform completions.Source: §1 · p. 2
  • The paper studies single-line completion, where the target hole can benefit from both surrounding-file context and relevant snippets from other repository files.The framework appends surrounding context to each repo context, encodes them separately, and combines their representations to predict the hole.Source: §1 · p. 2
  • A 220M-parameter RepoFusion model significantly outperforms larger next-token models such as CodeGen-16B and closely matches 15.5B-parameter StarCoderBase.The comparison with StarCoderBase uses a Fill-in-the-Middle objective; RepoFusion is approximately 70 times smaller.Source: §1 · p. 2
  • The authors release Stack-Repo, a dataset of 200 permissively licensed Java repositories with near-deduplicated files augmented with three types of repository context.The paper also reports extensive ablations on context sources, context lengths, numbers of contexts, and training configurations.Source: §1 · p. 2

2 Training with Repository Context

This section adapts Fusion-in-Decoder to code completion and describes repository contexts and strategies for turning prompt proposals into model inputs.

2.1 Fusion-in Decoder

  • Fusion-in-Decoder encodes each support passage independently, concatenates the representations, and lets the decoder jointly attend to them when producing an answer.The paper adapts this approach from open-domain question answering to code completion.Source: §2.1 · p. 2, Figure 1 · p. 3

2.2 Repository Contexts

  • Prompt proposals extract context types from related repository files; BM25 and RandomNN are also considered as repository-context retrieval mechanisms.Examples of prompt sources include imported or similarly named files, while context types include method bodies, identifiers, and string literals.Source: §2.2 · p. 3

2 Training with Repository Context

  • RepoFusion trains a code model to combine relevant information from multiple retrieved repository contexts to predict the target hole.It appends surrounding context to each repo context and encodes each pair independently.Source: §2.3 · p. 3

2.3 RepoFusion

  • The four strategies vary truncation and ordering: T-Rank truncates ranked proposals, T-Rand randomizes placement, NT-Rank uses full proposal chunks, and NT-Prior-Last puts prior context last.The number and length of repo contexts are N and l, respectively; NT-Prior-Last may remove top-ranked chunks to reserve room for prior context.Source: §2.3 · p. 4, Figure 2 · p. 4

3 Experiments and Results

The experiments evaluate RepoFusion on Java repository code completion against models of varying sizes and context setups. RepoFusion outperforms larger models in reported comparisons, while ablations show that context selection, ordering, length, number, and initialization affect performance.

3.1 Dataset Creation

  • Stack-Repo uses near-deduplicated Java repositories, split 2:1:1 into train, validation, and test sets, with target holes sampled from code lines and each repository capped at 10,000 holes.Repositories must contain at least 20 near-deduplicated files; target-hole positions are selected independently of any model tokenizer.Source: §3.1 · p. 5

3.2 Experimental Details

  • Evaluation compares fine-tuned encoder-decoder and decoder-only baselines using prior-only or post-plus-prior input, with success defined as an exact match to the target hole.The compared models include CodeT5, SantaCoder, StarCoderBase, and CodeGen variants; the maximum generation limit is 128 tokens per completion.Source: §3.2 · p. 6, §3.3 · p. 6

3.3 Results

  • RepoFusion surpasses larger models at the same effective context length and outperforms CodeGen-16B with fewer repository contexts; on a test-set subset, its success rate is 1.3% below StarCoderBase.Table 3 compares RepoFusion with StarCoderBase on a random subset of 12,500 test holes.Source: Table 3 · p. 7, §3.3 · p. 7
  • Prompt-proposal contexts perform best among the tested retrieval methods, while ordered strategies outperform random ordering and NT variants outperform truncated variants.The tested retrieval methods are prompt proposals, BM25, and Random-NN; NT-Prior-Last, NT-Rank, and T-Rank outperform T-Rand.Source: §3.3 · pp. 7–8, Figure 3 · p. 8
  • Increasing individual context length improves performance until it saturates for both tested context counts, N = 32 and N = 63.Figure 4 plots success rate against individual repository-context length.Source: §3.3 · p. 8, Figure 4 · p. 8
  • Performance improves with more repository contexts up to N = 63 at l = 512 or N = 32 at l = 768, then stops improving; diverse prompt proposals are important for better performance.Figure 4 also relates performance to the average number of prompt-proposal contexts producing the repository contexts.Source: §3.3 · p. 8, Figure 4 · p. 8
  • Across the tested cases, initializing RepoFusion from a CodeT5-base model fine-tuned for code completion benefits performance over initializing from the pretrained model.The comparison is shown in Table 4.Source: §3.3 · pp. 8–9, Table 4 · p. 9

4 Related Work

Prior work has used repository structure and retrieval to supply code models with context. RepoFusion builds on inference-time approaches by training a model to integrate multiple retrieved repository contexts.

4 Related Work

  • Repository structure has been modeled through locality-based caches, relational databases and graph walks, API-dependency graphs, and hierarchy features.Other work also uses parent classes to generate comments for child classes.Source: §4 · p. 9
  • Inference-time systems select prompt proposals or iteratively refine predictions using retrieved context, while RepoFusion trains models to leverage multiple repository contexts.RepoFusion builds on insights from RLPG [31] and RepoCoder [37].Source: §4 · p. 9
  • RepoFusion is independent of the retrieval mechanism and can integrate contexts from different mechanisms, including prompt proposals, BM25, and Random-NN.These approaches share similarities with systems that retrieve API documentation or relevant repository snippets.Source: §4 · p. 9

5 Discussion

The discussion identifies inference scalability as a limitation and summarizes RepoFusion’s single-line completion findings. It also points to Stack-Repo as a resource for further work.

5 Discussion

  • Inference computation scales linearly with the number of repo contexts N, making inference slower at larger N.The paper suggests FiDO [8] as one possible optimization; it also notes code-generation concerns common to code LLMs, including security risks and developer oversight of errors.Source: §5 · p. 9
  • RepoFusion’s single-line autocompletion experiments show performance enhancements from training smaller models with repository context compared with larger models without it.The authors suggest the framework may extend to bug repair, pull-request merging, and software documentation or tutorial writing.Source: §5 · pp. 9–10

A Details on Stack-Repo

  • Stack-Repo contains 200 near-deduplicated Java repositories, with repository files and PP, BM25, and RandomNN contexts provided for target holes.Each JSON row records the target-hole location and string, surrounding context, and repo contexts.Source: A Details on Stack-Repo · p. 13

B Implementation Details

The implementation details specify fine-tuning settings and compute used for CodeT5-base and CodeT5-large.

B Implementation Details

  • CodeT5-base (220M parameters) was fine-tuned with a 512-token input context, learning rate 4e-05, weight decay 0.05, and batch size 32 per GPU.Training used Adam with decoupled weight decay, 100 warm-up steps, dropout 0.1, and gradient clipping at 1.0.Source: §B.1 · p. 13
  • CodeT5-large (770M parameters) used a 512-token input context and the same hyperparameters as CodeT5-base except for a learning rate of 1e-4.Training used four 80GB NVIDIA A100 GPUs with batch size 12 per GPU.Source: §B.1 · p. 13
  • Evaluation used a single 32GB V100 GPU, with batch sizes of 32 for CodeT5-base and 48 for CodeT5-large.Source: §B.1 · p. 13

B.2 Training RepoFusion

RepoFusion uses CodeT5-base and FiD-based training with retrieved repository contexts; the appendix also records retrieval and baseline configurations.

B.2 Training RepoFusion

  • The RepoFusion implementation uses 220M-parameter CodeT5-base and builds on released code for prompt-proposal contexts and the FiD architecture.The cited implementations are from Shrivastava et al. [31] and Izacard and Grave [13].Source: §B.2 · p. 13
  • The best RepoFusion model starts from fine-tuned CodeT5-base and uses NT-Prior-Last with 32 prompt-proposal contexts of 768 tokens each.It was trained with learning rate 1e-5, 5,000 warm-up steps, gradient clipping at 1.0, and batch size 1; BM25 and Random-NN variants used N = 63 and l = 512 with learning rate 2.5e-5.Source: §B.2 · p. 13
  • The paper releases RepoFusion models and fine-tuned CodeT5 checkpoints.Source: §B.2 · p. 13

B.3 Retrieval Mechanisms

  • BM25 contexts use surrounding context as the query and repository-file contexts as search documents; Random-NN contexts use CodeBERT representations.BM25 uses the Okapi implementation with default rank-bm25 package parameters.Source: §B.3 · p. 14

B.4 Other Baselines

  • The baseline configurations include CodeGen-2B, 6B, and 16B, SantaCoder, and StarCoder, using FIM tokens and 128-token completion lengths.The models use their recommended tokenizers.Source: §B.4 · p. 14

C.1 Effect of Repetition

The ablation repeats prompt proposal contexts to test whether repetition from one proposal differs from using contexts from multiple proposals. Repeating a single proposal’s context hurts performance compared with using different repo contexts.

C.1 Effect of Repetition

  • Repeating PPCs from a single prompt proposal is compared with using different repo contexts from multiple proposals.The tested single proposals are prior, post, or randomly chosen.Source: §C.1 · p. 14, Table 5 · p. 14
  • Repeating context from one prior, post, or randomly chosen proposal negatively affects performance relative to using different repo contexts.The comparison is reported in the right side of Table 5.Source: §C.1 · p. 14, Table 5 · p. 14
  • Table 5 reports completion success rates for repeating different types of PPCs.Source: Table 5 · p. 14

C.2 Appending Surrounding Context

The ablation removes surrounding context from each repo context while using the best RepoFusion model trained on concatenated inputs. Performance drops significantly across all strategies, though the authors expect a smaller drop from a model trained without that context.

C.2 Appending Surrounding Context

  • Performance drops significantly across all strategies when surrounding context is removed from repo contexts.The comparison uses the best RepoFusion model, which was trained on concatenated surrounding and repo context.Source: §C.2 · p. 14, Table 6 · p. 14
  • The authors expect less performance loss if the model is trained without appended surrounding context.Source: §C.2 · p. 14
  • The ablation tests performance with and without appending surrounding context to each repo context.Table 6 reports completion success rates for the comparison.Source: Table 6 · p. 14, §C.2 · p. 14

C.3 Performance of Pretrained CodeT5

The pretrained CodeT5-base and CodeT5-large models perform poorly on the test set when prompted to complete the target hole. The authors report that finetuning on Java repositories with next-token prediction helps substantially.

C.3 Performance of Pretrained CodeT5

  • Pretrained CodeT5-base and CodeT5-large have quite low test-set performance when prompted with <extra_id_0> for target-hole completion.Table 7 reports completion success rates on the test set.Source: §C.3 · p. 15, Table 7 · p. 15
  • The authors conclude that pretrained CodeT5 needs finetuning on Java repositories with the next-token prediction objective.Source: §C.3 · p. 15
  • Finetuning helps substantially compared with directly using pretrained CodeT5 models.The text points to the top section of Table 2 in the main paper for this comparison.Source: §C.3 · p. 15
Loading 2306.10998v1…