CodeShrink

July 16, 2026 ยท View on GitHub

  • BFR/: Blank-Free Rendering, which renders source code into compact code images.
  • DTS/: Dominant Token Selection, which prunes visual tokens using foreground/background regions and attention scores.
  • ACC/: Adaptive Compression Configuration, which trains a configuration agent with SFT warm-up followed by GRPO.

1. Environment Setup

Python 3.10 or later is recommended.

python -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt

For BFR rendering only, the minimal dependencies are:

pip install pillow Pygments tiktoken tqdm

For ACC training, install the training dependencies:

pip install torch transformers ms-swift datasets

ACC uses Qwen3.5-0.8B as the configuration agent by default. A local model path can be specified with:

export MODEL_PATH=/path/to/Qwen3.5-0.8B

2. Data Preparation

Prepare the raw datasets following the public CodeOCR data preparation instructions, and keep only the three tasks used by this artifact:

  • Code Completion: download microsoft/LCC_python and microsoft/LCC_Java from Hugging Face.
  • Code QA: use qa_dataset_test_no_comments.json from the CodeOCR repository.
  • Code Clone Detection: prepare GPTCloneBench with the true_semantic_clones and false_semantic_clones directories.

Build task manifests

Code QA:

python main.py prepare-qa \
  --input /path/to/qa_dataset_test_no_comments.json \
  --output data/manifests/code_qa.jsonl

Code Completion:

python main.py prepare-completion \
  --dataset microsoft/LCC_python \
  --split test \
  --language python \
  --output data/manifests/completion_python.jsonl

python main.py prepare-completion \
  --dataset microsoft/LCC_Java \
  --split test \
  --language java \
  --output data/manifests/completion_java.jsonl

Code Clone Detection:

python main.py prepare-clone \
  --dataset /path/to/GPTCloneBench \
  --language python \
  --output data/manifests/clone_python.jsonl

python main.py prepare-clone \
  --dataset /path/to/GPTCloneBench \
  --language java \
  --output data/manifests/clone_java.jsonl

Render with BFR

python main.py render \
  --manifest data/manifests/code_qa.jsonl \
  --output-dir data/rendered/code_qa \
  --ratios 1,2,3,4,5,6,7,8

The render command writes images and a render_manifest.jsonl file containing sample IDs, rendering ratios, image paths, and image-token statistics.

Build ACC training data

ACC requires pre-evaluating each candidate configuration with the downstream model. Store the pre-evaluation results as environment.jsonl, where each row corresponds to one sample under one action:

{"sample_id":"sample-0001","task":"code_qa","language":"python","image":"data/rendered/code_qa/...png","action":{"r_delta":2,"r_phi_f":0.1,"r_phi_b":0.9},"compression":2.31,"correct":true}

Fields:

  • sample_id: sample identifier.
  • task: one of code_qa, code_completion, or code_clone_detection.
  • language: programming language.
  • image: image input for the configuration agent, usually the first page of the 1x BFR rendering.
  • action: compression configuration containing r_delta, r_phi_f, and r_phi_b.
  • compression: compression gain for this configuration; larger values indicate fewer retained visual tokens.
  • correct: whether the downstream model answers correctly under this configuration.

Generate SFT and GRPO data:

python main.py build-acc \
  --environment data/environment/environment.jsonl \
  --output-dir data/acc

Outputs:

data/acc/acc_sft.jsonl
data/acc/acc_grpo.jsonl
data/acc/acc_dataset_stats.json

3. How to Run

Generate the default discrete configuration space:

python main.py action-space --output data/action_space.json

Train the ACC agent after data/acc/acc_sft.jsonl and data/acc/acc_grpo.jsonl are prepared:

MODEL_PATH=/path/to/Qwen3.5-0.8B \
DATA_DIR=data/acc \
OUTPUT_DIR=outputs/acc \
bash ACC/train_ms_swift.sh

4. Quick Start

Run a minimal sanity check:

bash start.sh action-space

Full workflow:

python main.py prepare-qa/prepare-completion/prepare-clone ...
python main.py render ...
python main.py build-acc ...
bash start.sh train-acc