| # 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](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](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. |
|
|
