aiXcoder-colt

This repository contains the implementation of CoLT in our papers, a novel fine-tuning approach for enhancing large language models' ability to utilize information within long contexts for code completion tasks.

Resources

Dataset

• CoLT-132K dataset: A large-scale repo-level code completion dataset comprising 132,000 samples across four programming languages.
• Download: https://zenodo.org/records/15019938

Models

• Trained Models: This repository includes our models (aiXcoder-7B-v2, DS-Coder and Code Llama) trained with the CoLT approach.
• Download: https://zenodo.org/records/15021246

Experiment Reproduction

We use the latest TRL framework(https://github.com/huggingface/trl) code for our experiments. To reproduce our results:

1. Prompt Construction Files

• aiXcoder-colt/prompt/prompt_aixcoder_colt.py
• aiXcoder-colt/prompt/prompt_codellama.py
• aiXcoder-colt/prompt/prompt_codeqwen.py
• aiXcoder-colt/prompt/prompt_deepseekcoder.py

2. Training Scripts

• SFT (Supervised Fine-Tuning): See scripts in aiXcoder-colt/commands/sft/
• DPO (Direct Preference Optimization): See scripts in aiXcoder-colt/commands/po/

3. Reject Sampling for DPO

The aiXcoder-colt/Reject_Sample/ directory contains implementation and evaluation scripts for our reject sampling approach used in Direct Preference Optimization:

• Model-specific implementations:

  • aixcoder/: Reject sampling for aiXcoder model
  • codellama/: Reject sampling for Code Llama model
  • deepseek/: Reject sampling for DeepSeek-Coder model

• Evaluation scripts:

  • eval_api.py: API-based evaluation script
  • eval_line.py: Line-level evaluation script
  • eval_span.py: Span-level evaluation script
  • inference.py: Model inference script for generating completions

Dependencies

In our experiments, we utilized two Docker environments for TRL training and vLLM (reject sampling). Below are the key dependencies for each environment, excluding redundant packages:

TRL Training Environment:

• transformers==4.46.0.dev0
• torch==2.4.0a0+07cecf4168.nv24.5
• accelerate==1.0.0
• deepspeed==0.15.2
• peft==0.13.1
• flash-attn==2.4.2
• datasets==3.0.1
• wandb==0.15.0

vLLM Inference Environment:

• vllm==0.6.0+cu124
• torch==2.4.0
• transformers==4.44.2
• vllm-flash-attn==2.6.1
• xformers==0.0.27.post2
• flashinfer==0.1.6+cu121torch2.4
• fastapi==0.114.1
• uvicorn==0.30.6

Complete dependency lists can be found in the dependency directory.