CodeBERT Fine-Tuning for Code Search Task

This repository contains a PyTorch implementation of experiments for fine-tuning CodeBERT on the code search task, following the approach outlined in the paper CodeBERT: A Pretrained Model for Programming and Natural Languages (arXiv:2002.08155v4). The experiments are conducted using the CodeSearchNet dataset, and the primary evaluation metric is Mean Reciprocal Rank (MRR).

Workflow

1. Environment Setup
2. Dataset Download & Preprocessing
3. Model Fine-Tuning
4. Inference & Evaluation

Requirements

First, create and activate a virtual environment:

# Create and activate virtual environment
python -m venv env
source env/bin/activate  # On Windows use: env\Scripts\activate

Then, install the required dependencies:

pip install -r requirements.txt

Dataset setup

The dataset used for training and evaluation is the CodeSearchNet corpus, which includes code in multiple programming languages paired with natural language (NL) descriptions. For the fine-tuning process, both the training and validation datasets are balanced, with positive and negative samples. Negative samples are created by randomly replacing natural language (NL) and programming language (PL) pairs.

The evaluation is done using the Mean Reciprocal Rank (MRR) metric, computed over a fixed set of 999 distractor code snippets for each test pair (c, w).

Downloading and Preprocessing the Dataset

To prepare the dataset, follow these steps:

1. Create the necessary directories and navigate to the codesearch folder:

mkdir data data/codesearch
cd data/codesearch

2. Download and unzip the dataset:

gdown https://drive.google.com/uc?id=1xgSR34XO8xXZg4cZScDYj2eGerBE9iGo 
unzip codesearch_data.zip
rm  codesearch_data.zip

3. Run the preprocessing script :

cd ../../codesearch
python process_data.py
cd ..

This script processes the test data by dividing it into batches. For each batch, it creates one query–code pair where the code snippet is the correct match, along with 999 distractor pairs that act as negative examples.

Fine-Tuning

Fine-tuning a language specific model for each programming language.

lang=python  # java # javascript
pretrained_model=microsoft/codebert-base  # microsoft/codebert-base-mlm
model_name=${pretrained_model##*/}

python main.py \
--do_train \
--do_eval \
--eval_all_checkpoints \
--data_dir ./data/codesearch/train_valid/$lang \
--output_dir ./models/$model_name/$lang \
--model_name_or_path $pretrained_model
--train_file train.txt \
--dev_file valid.txt \
--overwrite_output_dir \  
[OPTIONS]

Training Parameters

Parameter	Default Value	Description
--train_batch_size	64	Batch size for training.
--eval_batch_size	64	Batch size for evaluation.
--max_seq_length	200	Max total input sequence length after tokenization.
--learning_rate	1e-5	The initial learning rate for the Adam optimizer.
--adam_epsilon	1e-8	Epsilon for the Adam optimizer.
--weight_decay	0.0	Weight decay for regularization.
--max_grad_norm	1.0	Maximum gradient norm for gradient clipping.
--num_train_epochs	8	Total number of training epochs.
--train_subset_ratio	1	Fraction of training data to use (between 0.0 and 1.0).
--early_stopping	0	Stop training if no improvement after this many epochs. "0" to disable.
--warmup_ratio	0.1	Linear warmup over warmup_ratio * total_steps.

Inference and Evaluation

Inference

lang=python # java # javascript
pretrained_model=microsoft/codebert-base-mlm  # microsoft/codebert-base
model_name=${pretrained_model##*/}
idx=0  # test batch idx


python main.py \
--model_name_or_path $pretrained_model \
--do_predict \
--test_file batch_${idx}.txt \
--output_dir ./models/$model_name/$lang  \
--data_dir ./data/codesearch/test/$lang \
--pred_model_dir ./models/$model_name/$lang/checkpoint-best/ \
--test_result_dir ./results/$model_name/$lang/${idx}_batch_result.txt \
[OPTIONS]

Test Parameters

Parameter	Default Value	Description
--max_seq_length	200	Max total input sequence length after tokenization.
--eval_batch_size	64	Batch size for evaluation.

Once inference is completed and prediction files are generated for all batches, it's possible to compute the final Mean Reciprocal Rank (MRR) score for each model and language using the script below:

Evaluation

python mrr.py

Experiment Results

All experimental results are tracked and available on wandb: View Results Dashboard