Algorithmic Inversion

File Structure

/
|——sources                              # source scripts
|  |—— AICoder.py                       # implementation of the algorithm embedding layer 
|  |—— tag_generator.py                 # implementation of the tag generator
|  |—— utils.py                         # helper functions
|  |—— AICoder_inference.py             # inference with the algorithm embedding layer
|  |—— tag_generator_inference.py       # generate topic tags
|  |—— train_AICoder.py                 # train the algorithm embedding layer
|  |—— train_tag_generator.py           # train the tag generator
|
|——eval                                 # evaluation scripts
|  |—— apps_metric                      # implementation of APPS evaluation from codeparrot
|  |—— code_contest_metric.py           # implementation of CodeContest evaluation
|  |—— eval_code_gen.py                 # run evaluation on APPS or CodeContest
|  |—— eval_tag_generator.py            # evaluate the tag generator
|
|——datasets                             # datasets for both training and evaluation
|  |—— APPS                             # APPS datasets
|  |—— CodeContest                      # CodeContest datasets
|  |—— training                         # training datasets
|
|...

Setup

Make sure that Python<=3.10 so as to be compatible with Pyext and Pytorch is properly installed.

Training of the algorithm embedding layer uses deepspeed, and the configuration can be found at /configs/deepspeed.

pip install -r requirements.txt

Training

Algorithm Embedding Layer

To run this training script, accelerate and deepspeed need to be properly configured according to your devices.

CUDA_VISIBLE_DEVICES="6,7" accelerate launch ./sources/train_AICoder.py \
    --input_path ./datasets/training/train-dedup-7005.jsonl \
    --output_path ./models/qwen/AICoder-67-ratio-codelen-prefix \
    --base_model_path ./models/qwen/Qwen2.5-Coder-1.5B  \
    --log_path ./logs/qwen/AICoder-67-ratio-codelen-prefix \
    --learning_rate 1e-3 \
    --max_epochs 10 \
    --per_device_batch_size 1 \
    --gradient_accumulation_steps 8 \
    --max_length 4096 \             # maximun input length
    --prefix_length 117 \           # picking from /configs/algorithm_embedding_configs
    --avg_loss_steps 10 \
    --save_every 10

Tag Generator

Training of the tag generator uses Pytorch DDP.

python ./sources/train_tag_generator.py \
    --input_path ./datasets/training/train-dedup-7005.jsonl \
    --output_path ./models/tag_generator/100-epochs \
    --base_model_path ./models/bge/bge-large-en-v1.5\
    --log_path ./logs/tag_generator/100-epochs \
    --learning_rate 1e-4 \
    --max_epochs 100 \
    --batch_size 256 \
    --devices 0,1

Inference

Algorithm Embedding Layer

This script has checkpoint mechanism which supports resuming from interruption like KeyboardInterrupt.

Batch size can be set high as this script will automatically reduce it if OOM happens.

python ./sources/AICoder_inference.py \
    --model_type qwen/AICoder-67-ratio-codelen-prefix \
    --metric CodeContest \      # select from "APPS" and "CodeContest"
    --split test \              # pick from /datasets
    --temperature 0.3 \
    --devices 3,4,5,6,7 \
    --batch_size 8 \
    --num_samples 5             # set k for pass@k

Tag Generator

python ./sources/tag_generator_inference.py \
    --model_type tag_generator/100-epochs\
    --output_path ./inference_results/tag-generator
    --device 7 \            # only support inference on single GPU
    --batch_size 256 \
    --threshold 0.97        # binary classification thresholds

Evaluation

python ./eval/eval_tag_generator.py --threshold 0.97