Repository files navigation

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.

• 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

• 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

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

• 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

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

TheaiXcoder-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

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:

• 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==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 thedependency directory.