PaperlayerView exact revision on arXiv

Source-linked AI summary

Teaching Code LLMs to Use Autocompletion Tools in Repository-Level Code Generation

Chong Wang, Jian Zhang, Yebo Feng, Tianlin Li, Weisong Sun, Yang Liu, Xin Peng

arXiv:2401.06391v3cs.SE

TL;DR

Repository-level code generation presents challenges in deciding when to invoke autocompletion and which suggestion to use when tools return multiple candidates. TOOLGEN trains code LLMs to invoke autocompletion during generation and select a suggestion; across three models, it improves dependency metrics while maintaining or improving similarity metrics.

  • Problem

    Integrating autocompletion requires identifying dependency-related trigger positions and selecting among multiple suggestions returned by the tool.

    Source: §I · p. 2

  • Method

    TOOLGEN fine-tunes a code LLM on functions marked with <COMP>, then invokes repository-context autocompletion at those markers and selects a suggestion during generation.

    Source: §I · pp. 2–3

  • Results

    Across three code LLMs, Dependency Coverage improved 31.4% to 39.1% and Static Validity Rate 44.9% to 57.7%, while similarity metrics remained comparable or improved.

    Source: §I · pp. 2–3

  • Takeaways & Limitations

    The evaluation reports high efficiency in repository-level generation, attributed to offline fine-tuning with trigger insertion.

    Source: §VIII · p. 14

    The implementation and evaluation are specific to Python, so findings may not generalize to other programming languages.

    Source: C. Threats to Validity · p. 14

Abstract

from arXiv · show

Code large language models (LLMs) face limitations in repository-level code generation due to their lack of awareness of repository-level dependencies (e.g., user-defined attributes), resulting in dependency errors such as undefined-variable and no-member errors. In this work, we introduce ToolGen, an approach that integrates autocompletion tools into the code LLM generation process to address these dependencies. ToolGen comprises two main phases: Trigger Insertion and Model Fine-tuning (Offline), and Tool-integrated Code Generation (Online). During the offline phase, ToolGen augments functions within a given code corpus with a special mark token, indicating positions to trigger autocompletion tools. These augmented functions, along with their corresponding docstrings, are then used to fine-tune a selected code LLM. In the online phase, ToolGen iteratively generates functions by predicting tokens step-by-step using the fine-tuned LLM. Whenever a mark token is encountered, ToolGen invokes the autocompletion tool to suggest code completions and selects the most appropriate one. We conduct comprehensive experiments to evaluate ToolGen's effectiveness in repository-level code generation. To facilitate this evaluation, we create a benchmark comprising 671 real-world code repositories and introduce two new dependency-based metrics: Dependency Coverage and Static Validity Rate. The results demonstrate that ToolGen significantly improves Dependency Coverage by 31.4% to 39.1% and Static Validity Rate by 44.9% to 57.7% across the three LLMs, while maintaining competitive or improved performance in widely recognized similarity metrics such as BLEU-4, CodeBLEU, Edit Similarity, and Exact Match. On the CoderEval dataset, ToolGen achieves improvements of 40.0% and 25.0% in Pass@1 for CodeT5 and CodeLlama, respectively.

I. INTRODUCTION

Repository-level code generation is difficult because code LLMs lack awareness of dependencies, while autocompletion tools can suggest contextually accessible identifiers. TOOLGEN teaches models when to invoke such tools and how to use their suggestions; across three code LLMs, it improves dependency metrics while maintaining or improving similarity metrics.

I. INTRODUCTION

  • Code LLMs often lack repository-level dependency awareness, causing undefined-variable and no-member errors that impede generated-code usability.The paper notes that more than 70% of real-world functions are not standalone.Source: §I · p. 1
  • TOOLGEN inserts <COMP> markers into training functions, fine-tunes a code LLM offline, then invokes autocompletion online and selects a suggestion with constraint greedy search.Markers are inserted at identifiers that autocompletion tools can recommend.Source: §I · pp. 2–3
  • The evaluation covers 12,406 Python functions from 671 repositories and 176 CoderEval tasks, using Dependency Coverage and Static Validity Rate alongside similarity metrics.The similarity metrics include BLEU-4, CodeBLEU, Edit Similarity, and Exact Match.Source: §I · p. 2
  • Across three code LLMs, Dependency Coverage rose 31.4% to 39.1% and Static Validity Rate 44.9% to 57.7%; CoderEval Pass@1 improved 40.0% for CodeT5 and 25.0% for CodeLlama.CodeGPT maintained the same Pass@1 rate on CoderEval.Source: §I · p. 3

A. Code LLMs

  • Decoder-only models predict from the description and prior tokens, whereas encoder-decoder models also condition on the encoder representation of the description.The paper identifies CodeGPT and CodeLlama as decoder-only examples and CodeT5 as an encoder-decoder example.Source: A. Code LLMs · p. 3

B. Autocompletion Tools

Autocompletion tools use repository context and a caret position to return code suggestions. In this paper, the relevant completions are identifiers, including variables, attributes, and functions.

B. Autocompletion Tools

  • An autocompletion tool takes a repository and caret position as input and returns completion suggestions.A caret position is specified by source file, line number, and column number.Source: B. Autocompletion Tools · pp. 3–4
  • The approach focuses on identifier-level completions, including partial identifiers, rather than keywords.The paper treats keywords as relatively straightforward for code LLMs to predict.Source: B. Autocompletion Tools · p. 4

III. APPROACH

  • TOOLGEN integrates autocompletion into code LLM generation to support repository-level code generation.Source: §III · p. 4

A. Overview

TOOLGEN has offline trigger insertion and model fine-tuning, followed by online generation that invokes autocompletion at learned positions. Trigger insertion marks eligible identifiers with <COMP> when the tool can suggest them.

A. Overview

  • TOOLGEN parses repository files into ASTs, extracts functions, inserts <COMP> at tool-suggested positions, pairs augmented functions with descriptions, and fine-tunes a code LLM to predict triggers.The description concatenates the function signature and docstring; functions without docstrings are omitted.Source: A. Overview · p. 4, 1) Trigger Insertion · p. 4

1) Trigger Insertion:

  • For each function body, trigger insertion checks non-built-in identifiers and inserts <COMP> before an identifier only when the completion tool can suggest it.The tool is queried at the identifier’s start position; these checks are intended to minimize unnecessary invocations.Source: 1) Trigger Insertion · p. 4
  • The method can apply beyond function bodies, but TOOLGEN focuses on functions because their docstrings provide descriptions more readily than code blocks outside functions.For code outside functions, extracting descriptions is challenging because the scope of line comments is difficult to determine.Source: 1) Trigger Insertion · p. 5

2) Model Fine-tuning:

TOOLGEN fine-tunes the code LLM on descriptions paired with augmented functions, then uses it with an autocompletion tool in iterative generation. A predicted <COMP> triggers suggestions that the LLM selects and appends.

2) Model Fine-tuning:

  • During fine-tuning, descriptions and augmented functions are tokenized and used for next-token generation, with cross-entropy loss against each ground-truth next token.For CodeLlama-7B, LoRA keeps most parameters fixed and adds trainable low-rank matrices to selected transformer layers.Source: 2) Model Fine-tuning · p. 5

1) Overall Process:

  • Generation iteratively predicts tokens from the description and incomplete function, appending each token until the LLM predicts <EOS>.The process starts with <BOS>; the LLM predicts a vocabulary-wide probability distribution and selects the highest-probability token.Source: 1) Overall Process · pp. 5–6
  • When the LLM predicts <COMP>, TOOLGEN requests repository-context completions, uses the LLM to select one, and appends the selected suggestion to the incomplete function.The inserted incomplete function has existing <COMP> tokens removed before the tool is invoked.Source: 1) Overall Process · p. 6

12 LLM-SELECT(D, F, C)

LLM-SELECT represents completion suggestions in a prefix tree and greedily follows a path guided by the fine-tuned LLM’s predictions. It restricts each next-token choice to tokens available among the current node’s children.

12 LLM-SELECT(D, F, C)

  • LLM-SELECT tokenizes each completion and inserts its token sequence into a prefix tree, where each root-to-terminal path represents a completion.The root stores the empty string, and terminal nodes mark the ends of token sequences.Source: 2) Completion Suggestion Selection · p. 6
  • At each tree node, LLM-SELECT masks the LLM’s next-token distribution to the tokens stored in that node’s children, then chooses the highest-probability remaining token.It repeats this process until reaching a terminal node, appending the selected tokens to the incomplete function.Source: 2) Completion Suggestion Selection · pp. 6–7
  • The selected path can supply a completion such as “_registered_updates,” whose tokens are appended to the incomplete function.In the Figure 5 example, those tokens lie along the green path.Source: 2) Completion Suggestion Selection · p. 7

IV. EVALUATION SETUP

The evaluation assesses TOOLGEN’s effectiveness and efficiency in repository-level code generation. Its research questions cover similarity, dependency coverage and errors, test passing, generation time, and generalizability across code LLMs.

A. Research Questions

  • The evaluation asks whether TOOLGEN’s generated code aligns with ground truth under common similarity metrics and covers repository dependencies while reducing dependency errors.The dependency question includes user-defined functions and attributes.Source: A. Research Questions · p. 7
  • The evaluation also asks how often generated functions pass tests, how long generation takes on average, and whether TOOLGEN works with different code LLMs.Source: A. Research Questions · p. 7

B. Implementation

The implementation uses Jedi for Python autocompletion and fine-tunes three code LLMs on functions augmented with completion triggers. Evaluation uses repository-derived CodeSearchNet data and CoderEval tasks.

B. Implementation

  • TOOLGEN is designed to be language-agnostic, but the prototype described here is Python-specific.Source: B. Implementation · p. 7
  • The implementation evaluates decoder-only CodeGPT and CodeLlama alongside encoder-decoder CodeT5, with model sizes ranging from 124 million to 7 billion parameters.Source: B. Implementation · p. 7
  • Jedi supplies completion suggestions that include repository-level dependencies such as user-defined attributes and functions.Source: B. Implementation · p. 7

1) Datasets:

  • The training corpus contains 249,298 augmented description-function pairs from 12,231 Python repositories; evaluation includes 12,406 functions from 671 repositories and 176 CoderEval tasks.Source: B. Implementation · p. 7, 1) Datasets · p. 8

B. Implementation

  • Fine-tuning uses full parameters for CodeGPT and CodeT5, but LoRA for CodeLlama, training 3.86% of its parameters.Source: B. Implementation · pp. 7–8

2) Baselines:

The study compares TOOLGEN with vanilla fine-tuned and retrieval-augmented baselines, using similarity, dependency-based, and execution-based evaluation metrics.

2) Baselines:

  • Vanilla baselines fine-tune the same base models on original functions, while RepoCoder variants add iterative retrieval-augmented generation.Source: 2) Baselines · p. 8
  • RAG-based TOOLGEN variants use the same retrieval process as the RepoCoder baselines, an approach described as orthogonal to tool integration.Source: 2) Baselines · p. 8

3) Metrics:

  • Evaluation combines similarity metrics, dependency metrics, and test pass rate; Dependency Coverage compares repository dependencies in ground truth with expressions in generated functions.Source: 3) Metrics · pp. 8–9
  • Static Validity Rate measures the proportion of generated functions passing static checks for syntax, no-member, and undefined-variable errors; Pass@1 measures functions passing all corresponding tests.Source: 3) Metrics · p. 9

V. RESULTS AND ANALYSES

TOOLGEN improves repository-level Dependency Coverage across all three base models, while similarity metrics are broadly competitive with mixed metric-specific changes. The authors note that trigger opportunities may be missed during long sequential generations.

A. RQ1: Similarity-based Effectiveness

1) Dependency Coverage:

2) Static Validity Rate:

TOOLGEN improves repository-level static validity and execution results over vanilla baselines, with reported generation latency of 0.63 to 2.34 seconds. Its trigger-token and caching design reduces unnecessary completion-tool calls.

2) Static Validity Rate:

  • ValRate rises 57.7%, 49.3%, and 44.9%, while ValRate-dep rises 67.7%, 60.0%, and 56.8% for CodeGPT, CodeT5, and CodeLlama, respectively.The analysis focuses on no-member and undefined-variable errors.Source: 2) Static Validity Rate · p. 11

C. RQ3: Execution-based Effectiveness

  • Against REPOCODER, TOOLGEN-gpt and TOOLGEN-t5 pass 1 and 2 more functions, respectively, while TOOLGEN-llama passes 15 versus REPOCODER-llama’s 19.The authors attribute REPOCODER-llama’s advantage over its vanilla baseline to its larger model and longer token limit benefiting retrieval-augmented prompts.Source: C. RQ3: Execution-based Effectiveness · p. 11
  • Adding RAG leaves overall pass rates unchanged, although ragTOOLGEN-llama gains 4 passing functions on project-runnable tasks.The results show different advantages and disadvantages of RAG integration across runnable-level dependencies.Source: C. RQ3: Execution-based Effectiveness · p. 11
  • Against the three VANILLA baselines, TOOLGEN generates 0, 2, and 3 additional functionally correct functions, corresponding to Pass@1 improvements of 0%, 40.0%, and 25.0%, respectively.Source: C. RQ3: Execution-based Effectiveness · p. 11

D. RQ4: Efficiency

  • Generation latency is 0.63 to 2.34 seconds per function, approximately twice the average generation time; trigger-token prediction and caching limit tool calls.The models predict an average of 5.02, 6.24, and 7.05 <COMP> tokens per task for CodeGPT, CodeT5, and CodeLlama, respectively.Source: D. RQ4: Efficiency · pp. 11–12

E. RQ5: Generalizability

TOOLGEN’s gains extend across model architectures and scales, but case studies and limitations identify boundaries in descriptions, static analysis, decoding, evaluation, and access to closed-source models.

E. RQ5: Generalizability

  • TOOLGEN improves or maintains dependency-based and execution-based metrics while remaining competitive on similarity metrics across model architectures and parameter scales.The evaluation reports consistent performance across architectures and scales, with an acceptable additional latency overhead.Source: E. RQ5: Generalizability · p. 12

A. Case Study

  • Case studies show TOOLGEN recovering Counter._value, but incomplete descriptions mislead both models and an earlier wrong token still derails dependency generation.In the third example, TOOLGEN predicts cls._get_service() but misses ServiceName.PLUGINS_MANAGER after choosing the wrong object for an argument.Source: A. Case Study · pp. 12–13

B. Limitations

  • The implementation is Python-specific, and static analysis may miss attributes when parameter types cannot be inferred.The authors propose exploring comprehensive type-inference tools alongside autocompletion.Source: B. Limitations · p. 13
  • Greedy next-token prediction can select sub-optimal tokens and produce lower-quality code; beam search and other decoding methods remain future work.Source: B. Limitations · p. 13
  • Static analysis may make Dependency Coverage and Static Validity Rate inaccurate, although the same analysis is used for baselines.The authors state this does not significantly impact the demonstrated effectiveness of TOOLGEN.Source: B. Limitations · p. 13
  • The evaluation excludes GPT-3.5 and GPT-4 because their closed-source nature makes trigger insertion and fine-tuning impractical.The comparison therefore focuses on TOOLGEN, VANILLA, and REPOCODER under the same base models.Source: B. Limitations · pp. 13–14

C. Threats to Validity

The paper identifies validity threats from data quality, possible CodeLlama data leakage, and Python-only evaluation. It presents autocompletion integration as a response to missing repository dependencies and reports improvements across several metrics and models.

C. Threats to Validity

  • The benchmark may contain data quality issues and repositories encountered during CodeLlama pretraining, but consistent RQ5 performance suggests TOOLGEN-llama’s gains are not attributed to leakage.The dataset and benchmark are constructed using CodeSearchNet to mitigate data-quality concerns.Source: C. Threats to Validity · p. 14
  • The implementation and evaluation are specific to Python, so the findings may not generalize to other programming languages.Source: C. Threats to Validity · p. 14

A. Code Generation

  • Repository-level generation is challenged by missing user-defined dependencies; TOOLGEN integrates autocompletion tools into code LLM generation to address that challenge.Source: A. Code Generation · p. 14

B. Incorporating External Tools into LLMs

  • Prior tool integrations include arithmetic and API search [31], but API search does not address repository-level dependencies; TOOLGEN integrates program-analysis-based autocompletion.The cited approaches and comparisons are described in the paper’s related-work discussion.Source: B. Incorporating External Tools into… · p. 14

VIII. CONCLUSION

  • TOOLGEN improves two dependency-level metrics and an execution-based metric across three code LLMs, remains competitive on similarity metrics, and shows efficiency and generalizability across model architectures and scales.Source: §VIII · p. 14
Loading 2401.06391v3…