ReLOpt: A Retriever-Augmented Framework for Optimizing Code with Long-range Dependencies

This project provides the code for reproducing the experiments in ReLOpt: A Retriever-Augmented Framework for Optimizing Code with Long-range Dependencies.

ReLOpt is a novel framework designed to augment the capabilities of LLMs for code optimization with long-range dependencies.

1. Fine-Tuning

You can fine-tune ReLOpt on our datasets by running (you need change the environment variables in the script to the right path):

bash ./Src/Model/run_fine_tuning.sh

You need to specify following options.

parser.add_argument('-m','--model')
parser.add_argument('-token-num','--token_num')
parser.add_argument('-lr','--lr')
parser.add_argument('-dataset-dir', '--dataset_dir')
parser.add_argument('-epochs', '--epochs')
parser.add_argument('-output-dir', '--output_dir')
parser.add_argument('-adapters-dir', '--adapters_dir', 
					default=None, 
					help="If you want to resume from checkpoint, please give the adapter dir.")

Fined-tuned models will be saved in the output-dir you specify.

2. Inference

You can use our model to inference by running (you need change the environment variables in the script to the right path):

bash ./Src/Model/run_infer.sh

You need to specify following options.

parser.add_argument('-m','--model')
parser.add_argument('-token-num','--token_num')
parser.add_argument('-infer-kind','--infer_kind', 
					required=True, help="'ReLOpt' or 'ReLOpt-off-retriever' or 'Unopt2Opt'.")
parser.add_argument('-test-dataset-dir', '--test_dataset_dir')
parser.add_argument('-adapters-dir', '--adapters_dir')
parser.add_argument('-output-dir', '--output_dir')

Inference results will be saved in the output-dir you specify.

3. Reproducing Experiment Results

We provide code generated by ReLOpt in ./Src/Model/ModelRes/ReLOpt.json.

We provide code generated by Code Llama - Instruct (7B) and CodeGemma (7B) in ./Src/Model/ModelRes/BaseModels.json.

The metrics used in the paper can be reproduced by running:

bash ./Src/Model/reproduce_metrics.sh