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_pythonandmicrosoft/LCC_Javafrom Hugging Face. - Code QA: use
qa_dataset_test_no_comments.jsonfrom the CodeOCR repository. - Code Clone Detection: prepare GPTCloneBench with the
true_semantic_clonesandfalse_semantic_clonesdirectories.
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 ofcode_qa,code_completion, orcode_clone_detection.language: programming language.image: image input for the configuration agent, usually the first page of the 1x BFR rendering.action: compression configuration containingr_delta,r_phi_f, andr_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