Skip to content
This repository has been archived by the owner on Jul 30, 2024. It is now read-only.
/ PLBART Public archive

Official code of our work, Unified Pre-training for Program Understanding and Generation [NAACL 2021].

License

Notifications You must be signed in to change notification settings

wasiahmad/PLBART

Repository files navigation

PLBART

Official code release of our NAACL 2021 work, Unified Pre-training for Program Understanding and Generation .

***** PLBART's performances on the downstream tasks are recorded in this spreadsheet. *****

NewsSetupPre-trainingFine-tuningFAQsAcknowledgementLicenseCitation


PLBART is a Transformer model

  • PLBART is a sequence-to-sequence model pre-trained on a large collection Java and Python functions and natural language descriptions collected from Github and StackOverflow, respectively.
  • PLBART is pre-trained via denoising autoencoding (DAE) and uses three noising strategies: token masking, token deletion, and token infilling (shown below in the three examples).
Noisy Input Original Sequence
Is 0 the [MASK] Fibonacci [MASK] ? <En> <En> Is 0 the first Fibonacci number ?
public static main ( String args [ ] ) { date = Date ( ) ; System . out . ( String . format ( " Current Date : % tc " , ) ) ; } <java> <java> public static void main ( String args [ ] ) { Date date = new Date ( ) ; System . out . printf ( String . format ( " Current Date : % tc " , date ) ) ; }
def addThreeNumbers ( x , y , z ) : NEW_LINE INDENT return [MASK] <python> <python> def addThreeNumbers ( x , y , z ) : NEW_LINE INDENT return x + y + z

News


Setup

We can setup a conda environment in order to run PLBART experiments, the first step is to download the dependencies. We assume anaconda is installed. The additional requirements (noted in requirements.txt) can be installed by running the following script:

bash install_env.sh

Pre-training

Step1. Download Github data

Go to data/github directory and follow instructions.

Step2. Download StackOverflow data

Go to data/stackoverflow directory and follow instructions.

Step3. Binarize the data and pre-train

cd pretrain
bash binarize.sh
bash pretrain.sh GPU_IDS

[Note] We pre-trained PLBART on 8 GeForce RTX 2080 (11gb) GPUs (took ~11.5 days). If you want to pre-train PLBART using more GPUs or GPUs with more memory, adjust MAX_SENTENCES, MAX_TOKENS, UPDATE_FREQ accordingly to maintain an effective batch size of 2048. According to fairseq, effective batch size is equal to:

PER_GPU_TRAIN_BATCH_SIZE * NUM_GPU * UPDATE_FREQ

Note that, MAX_TOKENS refers to the size of each mini-batch, in terms of the number of tokens. During our experiments, we noticed that in an 11gb GPU, maximum 2048 tokens can be accommodated which is equivalent to 4-5 examples. Therefore, we set UPDATE_FREQ to 60, so that we can achieve an effective batch size of ~2048.


Fine-tuning

We fine-tune and evaluate PLBART on three types of downstream tasks.

Type Task Language(s) Data Scripts Checkpoints
Code to Text Code summarization Python, Java, Ruby,
PHP, Javascript, Go
[LINK] [LINK] [LINK]
Text to Code Code generation Java [LINK] [LINK] [LINK]
Code to Code Code translation Java, C# [LINK] [LINK] [LINK]
Code refinement Java [LINK] [LINK]
Clone detection Java [LINK] [LINK]
Defect detection C/C++ [LINK] [LINK]

Step1. Download PLBART checkpoint

cd pretrain
bash download.sh
cd ..

Step2. Download the data

cd data/codeXglue
bash download.sh
cd ../..

Step3. Build parser for CodeBLEU evaluation

cd evaluation/CodeBLEU/parser
bash build.sh
cd ../../..

Step4. Prepare the data, train and evaluate PLBART

For example, we want to fine-tune PLBART on Text-to-Code task. Then,

cd scripts/text_to_code
bash prepare.sh
bash run.sh GPU_IDS
cd ../..

Note. We fine-tuned PLBART on 1 GeForce RTX 2080 (11gb) GPU.


FAQs

[NOTE] We present the file structure of this repository here .

How to download Github data from Google BigQuery?

We provided a detailed guide here .

Mismatch in performance reported in the paper and achieved using the released checkpoints.

There is a difference between PLBART's performances mentioned in the paper and the performance achieved with the released checkpoints. We noted them here. Note that, there is no change in the hyper-parameter setting. We provided the exact same value we used in the bash scripts. The performance difference we observed is perhaps due to running experiments at different point of time. Although we didn't but we recommend to fine-tune PLBART with multiple different seeds and report the average scores.

mbart_base task is not present in fairseq==0.9.0 official release.

Although we used fairseq==0.9.0 but we used a different commit which consists of mbart_base task. You may do the following which should work.

git clone https://github.com/pytorch/fairseq
cd fairseq
git checkout 698e3b91ffa832c286c48035bdff78238b0de8ae
pip install .

Otherwise, you may consider installing fairseq==0.10.0. Please refer to this issue to make other adjustments.

What can be the maximum input and output lengths for PLBART?

The maximum length is 512.


Acknowledgement

PLBART uses Fairseq, codeXglue, and TransCoder and thanks the authors of these works for their contribution.


License

Contents of this repository is under the MIT license. The license applies to the pre-trained and fine-tuned models as well.


Citation

@inproceedings{ahmad-etal-2021-unified,
    title = "Unified Pre-training for Program Understanding and Generation",
    author = "Ahmad, Wasi  and
      Chakraborty, Saikat  and
      Ray, Baishakhi  and
      Chang, Kai-Wei",
    booktitle = "Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies",
    month = jun,
    year = "2021",
    address = "Online",
    publisher = "Association for Computational Linguistics",
    url = "https://www.aclweb.org/anthology/2021.naacl-main.211",
    pages = "2655--2668"
}