Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

CODEPROMPTZIP: Code-specific Prompt Compression for Retrieval-Augmented Generation in Coding Tasks with LMs

Pengfei He, Shaowei Wang, Tse-Hsun Chen

TL;DR

CodePromptZip tackles the challenge of long prompts in retrieval-augmented generation for coding tasks by training a code-specific compressor with a type-aware token-priority ranking and a copy-enhanced CodeT5 backbone. The method enables ratio-controlled compression of retrieved code before inclusion in prompts, preserving generation quality while reducing token cost. Empirical results show meaningful gains over state-of-the-art NL-oriented prompt compression baselines across three coding tasks, with demonstrated transferability across base LMs and robustness to unparsable code. The work offers a practical path to cost-efficient, scalable RAG in coding contexts and supplies a general framework for code-aware prompt compression.

Abstract

Retrieval-Augmented Generation (RAG) enhances coding tasks by incorporating retrieved code examples into prompts. However, lengthy prompts, often exceeding tens of thousands of tokens, introduce challenges related to limited context windows of language models (LMs) and high computational costs. Existing prompt compression techniques focus on natural language, lacking tailored solutions for code. To address the gap, we propose CodePromptZip, a framework that compresses code examples before integrating into RAG workflows. Our framework employs a type-aware, priority-driven strategy to construct training samples for training code compression model. By using program analysis, we identify token types (e.g., Identifier) and perform ablation analysis to rank their removal priorities based on their impact on task performance. We then train a small LM as the compressor on these samples, enabling flexible compression conditioned on specified ratios while minimizing performance degradation. Specially, the compressor is augmented with a copy mechanism, allowing tokens to be directly copied from the original code snippets. Evaluation results show that CodePromptZip surpasses SOTA entropy-based and distillation-based baselines, improving by 23.4%, 28.7%, and 8.7% over the best baseline for Assertion Generation, Bugs2Fix, and Code Suggestion, respectively.

CODEPROMPTZIP: Code-specific Prompt Compression for Retrieval-Augmented Generation in Coding Tasks with LMs

TL;DR

CodePromptZip tackles the challenge of long prompts in retrieval-augmented generation for coding tasks by training a code-specific compressor with a type-aware token-priority ranking and a copy-enhanced CodeT5 backbone. The method enables ratio-controlled compression of retrieved code before inclusion in prompts, preserving generation quality while reducing token cost. Empirical results show meaningful gains over state-of-the-art NL-oriented prompt compression baselines across three coding tasks, with demonstrated transferability across base LMs and robustness to unparsable code. The work offers a practical path to cost-efficient, scalable RAG in coding contexts and supplies a general framework for code-aware prompt compression.

Abstract

Retrieval-Augmented Generation (RAG) enhances coding tasks by incorporating retrieved code examples into prompts. However, lengthy prompts, often exceeding tens of thousands of tokens, introduce challenges related to limited context windows of language models (LMs) and high computational costs. Existing prompt compression techniques focus on natural language, lacking tailored solutions for code. To address the gap, we propose CodePromptZip, a framework that compresses code examples before integrating into RAG workflows. Our framework employs a type-aware, priority-driven strategy to construct training samples for training code compression model. By using program analysis, we identify token types (e.g., Identifier) and perform ablation analysis to rank their removal priorities based on their impact on task performance. We then train a small LM as the compressor on these samples, enabling flexible compression conditioned on specified ratios while minimizing performance degradation. Specially, the compressor is augmented with a copy mechanism, allowing tokens to be directly copied from the original code snippets. Evaluation results show that CodePromptZip surpasses SOTA entropy-based and distillation-based baselines, improving by 23.4%, 28.7%, and 8.7% over the best baseline for Assertion Generation, Bugs2Fix, and Code Suggestion, respectively.

Paper Structure

This paper contains 32 sections, 10 equations, 14 figures, 6 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work and Background
  3. Related work
  4. Problem Formulation
  5. Type-aware Priority Ranking
  6. Type Ablation Analysis
  7. Setup of Downstream Coding Tasks with RAG
  8. Observation
  9. Methodology
  10. Code Compression Dataset Construction
  11. Compressor Architecture
  12. Experimental Setting
  13. Research Questions
  14. Baselines and Oracle
  15. Datasets and Metrics
  16. ...and 17 more sections

Figures (14)

  • Figure 1: Removal priority of code token types: e.g., Invocation > Symbol in Assertion Generation, and vice versa in Code Suggestion. Priorities are task-specific yet model-agnostic, applicable to both LMs.
  • Figure 2: Framework of CodePromptZip.
  • Figure 3: Illustration of copy mechanism on CodeT5.
  • Figure 4: Trade-off between keeping more tokens in a single example or including more examples.
  • Figure 5: Compression ratio control.
  • ...and 9 more figures