RepoFusion: Treinando Modelos de Código para Compreender Seu Repositório

RepoFusion: Training Code Models to Understand Your Repository

June 19, 2023

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

cs.AI

Resumo

Apesar do enorme sucesso dos Modelos de Linguagem de Grande Escala (LLMs) em assistentes de codificação como o GitHub Copilot, esses modelos têm dificuldade em compreender o contexto presente no repositório (por exemplo, imports, classes pai, arquivos com nomes semelhantes, etc.), resultando em sugestões de código imprecisas. Esse efeito é mais pronunciado ao usar esses assistentes para repositórios que o modelo não viu durante o treinamento, como software proprietário ou projetos de código em andamento. Trabalhos recentes mostraram o potencial de usar o contexto do repositório durante a inferência. Neste trabalho, estendemos essa ideia e propomos o RepoFusion, um framework para treinar modelos a incorporar o contexto relevante do repositório. Experimentos em conclusão de código de linha única mostram que nossos modelos treinados com contexto de repositório superam significativamente modelos de código muito maiores, como o CodeGen-16B-multi (aproximadamente 73 vezes maior), e se aproximam do desempenho do modelo StarCoderBase, que é aproximadamente 70 vezes maior e foi treinado com o objetivo Fill-in-the-Middle. Consideramos esses resultados uma demonstração nova e convincente dos ganhos que o treinamento com contexto de repositório pode trazer. Realizamos extensos estudos de ablação para investigar o impacto de escolhas de design, como tipo de contexto, número de contextos, comprimento do contexto e inicialização dentro de nosso framework. Por fim, lançamos o Stack-Repo, um conjunto de dados de 200 repositórios Java com licenças permissivas e arquivos quase deduplicados, que são aumentados com três tipos de contextos de repositório. Além disso, estamos disponibilizando o código e os checkpoints treinados para nosso trabalho. Nossos recursos lançados podem ser encontrados em https://huggingface.co/RepoFusion.

English

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 (sim73times larger) and closely match the performance of the sim 70times 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 https://huggingface.co/RepoFusion.