Code-Specific Prompt Compression for
Retrieval-Augmented Generation in Bug Fixing

Compress code demonstrations by up to 87% with near-zero quality loss

A reproduction and extension of the CodePromptZip framework (He et al., 2025), implementing a complete pipeline for code-specific prompt compression applied to the Bugs2Fix task from the Microsoft CodeXGLUE benchmark. The system compresses retrieved code demonstrations by up to 87% while preserving downstream bug-fixing quality, using a Copy-Enhanced CodeT5 compressor and CodeLlama-34B as the base language model.

Authors: Ansh Gupta (IMT2023540) · Mayank Tanwar (IMT2023012) · Prakrititz Borah (IMT2023547) · Kunal Jindal (IMT2023049)

Course: NLP + Deep Learning Project, IIIT Bangalore, April 2026

📑 Table of Contents

🎯 Problem Statement

Modern LLM-based code repair pipelines use Retrieval-Augmented Generation (RAG): before asking the model to fix a bug, similar buggy–fixed code pairs are retrieved from a knowledge base and placed into the prompt as demonstrations. While effective, this creates two critical bottlenecks:

1. Context window limits — Retrieved examples containing full buggy and fixed methods consume thousands of tokens, often exceeding the LLM's context window.

2. High inference cost — Proprietary APIs charge per token, making bloated prompts prohibitively expensive.

Existing prompt compression methods (LLMLingua, RECOMP, LLMLingua-2) were designed for natural language and fail on code because they do not consider the syntactic role of tokens. An entropy-based approach might discard a method invocation like .getInstance() because its sub-tokens are predictable, even though it is semantically critical for understanding a bug.

CodePromptZip addresses this gap by introducing a code-specific prompt compression framework that leverages program analysis to understand which tokens are safe to remove for a given downstream task.

🏗️ Architecture Overview

The framework operates in two fully decoupled phases — the compressor and the base LM share no gradients or embeddings.

Training Phase

╭─────────────────╮     ╭──────────────────╮     ╭──────────────────╮     ╭───────────────────╮     ╭─────────────────────╮
│  Knowledge Base │     │ Program Analysis │     │ Ablation Analysis│     │  Priority-Driven  │     │   Fine-tune         │
│  52,364 pairs   │────▸│ (JavaParser/AST) │────▸│ (Token Priority  │────▸│  Greedy Algorithm │────▸│   Copy-Enhanced     │
│  (Bugs2Fix)     │     │                  │     │  Ranking)        │     │  (9 ratios)       │     │   CodeT5-large      │
╰─────────────────╯     ╰──────────────────╯     ╰──────────────────╯     ╰───────────────────╯     ╰─────────────────────╯

Inference Phase

╭──────────────╮     ╭───────────────╮     ╭────────────────╮     ╭──────────────────╮     ╭──────────────╮
│  New Buggy   │     │ BM25 Retriever│     │    LM_C        │     │    Base LM       │     │  Generated   │
│  Query       │────▸│ (top-k similar│────▸│  Compressor    │────▸│  CodeLlama-34B   │────▸│  Fix         │
│              │     │  pairs)       │     │  (target τ)    │     │  (Q4 quantized)  │     │              │
╰──────────────╯     ╰───────────────╯     ╰────────────────╯     ╰──────────────────╯     ╰──────────────╯

🔧 Key Technical Components

1. Type-Aware Priority Ranking

Using JavaParser to construct Abstract Syntax Trees (ASTs), every token is categorized into one of five types:

The removal priority is computed as:

Priority(T) = τ_code/T / d_T

where τ_code/T is the compression ratio achieved by removing all tokens of type T, and d_T is the resulting CodeBLEU degradation.

Bugs2Fix Priority Ranking:

🏷️ Identifier  ▸  📞 Invocation  ▸  🔀 Structure  ▸  🔣 Symbol  ▸  📝 Signature
  (remove first)                                                      (remove last)

2. Priority-Driven Greedy Compression

For each parsable code example, the greedy algorithm:

1. Assigns each token a removal priority based on its type and within-type frequency
2. Computes L_rm = ⌊τ_code × L⌋ — the number of tokens to remove
3. Iteratively removes the highest-priority tokens until the budget is exhausted
4. Returns the compressed code as the training target

This is repeated for 9 compression ratios (τ_code ∈ {0.1, 0.2, ..., 0.9}):

⚠️ The split is performed at the code-example level (not at the triple level) to prevent data leakage — all 9 compressed versions of the same code example are assigned to the same split.

3. Copy-Enhanced CodeT5 Compressor

The compressor is built on CodeT5-large (Salesforce, 770M parameters), an encoder-decoder transformer pre-trained on 8.35M code functions across 8 programming languages.

<BUGS2FIX> <Ratio> 0.3 </Ratio> <Compress> {code} </Compress>

P(y) = p_gen · P_vocab(y) + (1 - p_gen) · P_copy(y)

📂 Project Structure

CodePromptZip/
│
├── 📁 configs/
│   └── config.yaml                 # All hyperparameters and settings
│
├── 📁 data/
│   └── download_datasets.py        # Download Bugs2Fix from HuggingFace
│
├── 📁 src/
│   ├── 📁 model/
│   │   ├── copy_codet5.py          # CopyCodeT5: CodeT5 + Copy Mechanism
│   │   └── copy_module.py          # Pointer-Generator Copy Module
│   ├── 📁 metrics/
│   │   ├── codebleu_metric.py      # CodeBLEU metric wrapper
│   │   └── exact_match.py          # Exact Match metric
│   ├── type_analysis.py            # JavaParser AST token categorization
│   ├── priority_ranking.py         # Priority-driven greedy algorithm
│   ├── dataset_construction.py     # Multi-ratio compression dataset builder
│   ├── tokenizer_utils.py          # Extended tokenizer with special tokens
│   ├── compress.py                 # Neural compressor inference
│   ├── retrieval.py                # BM25 retriever + RAG prompt formatting
│   ├── train.py                    # Training loop (FP16 + grad checkpointing)
│   └── evaluate.py                 # Full evaluation pipeline
│
├── 📁 scripts/
│   ├── run_all_evaluations.py      # Batch eval across all configurations
│   ├── plot_results.py             # Generate result plots and tables
│   ├── run_quick_test.py           # Quick sanity check
│   ├── run_train_compressor.py     # Training launcher
│   └── run_compress_and_eval.py    # Compression + eval launcher
│
├── 📁 results/
│   └── 📁 plots/                   # Generated visualizations
│
├── demo_viva.py                    # Interactive step-by-step demo
├── setup_env.sh                    # One-command environment setup (Linux)
├── requirements.txt                # Python dependencies
└── EXECUTION_FLOW.md               # Detailed execution guide

⚡ Setup & Installation

Prerequisites

Quick Start

# 1. Clone the repository
git clone https://github.com/prakrititz/Token-PruningNLP.git
cd Token-PruningNLP

# 2. Run the automated setup script
chmod +x setup_env.sh
bash setup_env.sh

# 3. Activate the virtual environment
source venv/bin/activate

# 4. Download the Bugs2Fix dataset
python data/download_datasets.py --task bugs2fix

# 5. Verify GPU is detected
python -c "import torch; print(f'CUDA: {torch.cuda.is_available()}, GPU: {torch.cuda.get_device_name(0)}')"

🚀 Execution Flow

Step 1: Download Dataset

python data/download_datasets.py --task bugs2fix

Downloads 52,364 train / 6,546 val / 6,545 test buggy–fixed Java code pairs into data/raw/bugs2fix/.

Step 2: Construct Compression Dataset

python src/dataset_construction.py --config configs/config.yaml

Produces 471,276 compression triples (52,364 × 9 ratios) split into train/val/test at the example level → data/compression_dataset/.

Step 3: Train the Compressor

python src/train.py --config configs/config.yaml

The best checkpoint (lowest validation loss) is auto-saved to checkpoints/best_model/.

Step 4: Run Evaluation Suite

# Full evaluation (2000 samples × 13 experiments)
python scripts/run_all_evaluations.py \
    --config configs/config.yaml \
    --checkpoint ./checkpoints/best_model \
    --max_eval_samples 2000

# Resume after interruption
python scripts/run_all_evaluations.py \
    --config configs/config.yaml \
    --checkpoint ./checkpoints/best_model \
    --max_eval_samples 2000 \
    --skip_existing

Step 5: Plot Results

python scripts/plot_results.py

Step 6: Interactive Demo

python demo_viva.py                  # Full pipeline with CodeLlama
python demo_viva.py --no_codellama   # Skip CodeLlama (retrieval + compression only)
python demo_viva.py --no_compressor  # Use greedy algorithm instead of neural compressor

⚙️ Training Configuration

All hyperparameters are defined in configs/config.yaml.

Training Curve

CopyCodeT5 Training & Validation Loss across 10 epochs. Best validation loss: 0.0178 (achieved at ~epoch 2).

📊 Results

All results below are from 2,000 evaluation samples (seed=42) using the trained CodeT5-large compressor and CodeLlama-34B-Instruct (Q4_K_M quantized) as the base LM.

Experiment 1: Compression Ratio Sweep (1-shot)

Sweep τ_code from 0.0 (no compression) to 0.9 (extreme compression) with num_shots = 1.

CodeBLEU remains remarkably stable across all compression levels. Even at 87% compression (τ=0.9), quality is virtually identical to the uncompressed baseline.

Experiment 2: Multi-Shot Ablation (τ = 0.3)

Fix τ_code = 0.3 and vary the number of retrieved demonstrations.

With compression (blue), performance degrades gracefully from 1-shot to 3-shot. Without compression (gray), it slightly improves.

Experiment 3: Uncompressed Baselines

Evaluate uncompressed demonstrations at multiple shot counts to isolate the effect of compression vs. retrieval.

Token Savings Visualization

Original vs. compressed token counts across configurations. At τ=0.9, demonstrations shrink from 160 to just 20 tokens — an 88% reduction.

Complete Ablation Summary

All 13 experimental configurations ranked by CodeBLEU. The narrow spread (87.1–88.7) demonstrates the robustness of both the compressor and CodeLlama-34B.

🔑 Key Findings

1. Non-monotonic compression curve — CodeBLEU does not degrade linearly with compression. After an initial dip at low τ (0.1–0.3), performance recovers at higher compression (0.5–0.9). At τ = 0.5, the compressor achieves 40.9% token reduction with a 0.08% CodeBLEU improvement over the uncompressed baseline (88.38 vs. 88.30).

2. Extreme compression robustness — Even at τ = 0.9 (87% token reduction, 160 → 20 tokens), CodeBLEU remains at 88.39 — essentially identical to the uncompressed baseline. This demonstrates that CodeLlama-34B is highly robust to compressed demonstrations.

3. τ = 0.1 expansion anomaly — At τ = 0.1, the compressor actually increases the token count (160 → 167). The CodeT5 decoder generates tokens auto-regressively, and at very low compression targets, it may reproduce the input with minor additions.

4. Graceful multi-shot degradation — With compression at τ = 0.3, performance degrades gracefully from 1-shot (87.89) to 3-shot (87.33), a mild 0.56% drop. Without compression, performance slightly improves with more shots (88.31 → 88.74).

5. Strong zero-shot baseline — Zero-shot achieves 88.10 CodeBLEU, indicating CodeLlama-34B is already highly capable at bug fixing without demonstrations.

🔬 Differences from the Original Paper

🤗 Pretrained Model

The trained Copy-Enhanced CodeT5-large compressor checkpoint is available on HuggingFace:

🤗 HuggingFace Model: CodePromptZip-Token-Pruning

Quick usage:

from src.compress import CodeCompressor

compressor = CodeCompressor("./checkpoints/best_model")  # or HuggingFace model ID
compressed = compressor.compress(
    "public int add(int a, int b) { return a - b; }",
    tau_code=0.3,
    task="bugs2fix"
)
print(compressed)

💻 Smaller Scale Implementation

A related, smaller-scale implementation of this project (trained and evaluated on an NVIDIA RTX 4060) can be found here:

🔗 GitHub Repository: prakrititz/Token-PruningNLP

📚 References