CodeT5-large was pretrained using Masked Span Prediction (MSP) objective on CodeSearchNet and achieve new SOTA results on several CodeXGLUE benchmarks. The finetuned checkpoints are released at here. See Appendix A.1 of the paper for more details.
CodeT5-large-ntp-py was first pretrained using Masked Span Prediction (MSP) objective on CodeSearchNet and GCPY (the Python split of Github Code data), followed by another 10 epochs on GCPY using Next Token Prediction (NTP) objective.
CodeT5-large-ntp-py is especially optimized for Python code generation tasks and employed as the foundation model for our CodeRL, yielding new SOTA results on the APPS Python competition-level program synthesis benchmark. See the paper for more details.
1from transformers import RobertaTokenizer, T5ForConditionalGeneration
23tokenizer = RobertaTokenizer.from_pretrained('Salesforce/codet5-base')4model = T5ForConditionalGeneration.from_pretrained('Salesforce/codet5-base')56text ="def greet(user): print(f'hello <extra_id_0>!')"7input_ids = tokenizer(text, return_tensors="pt").input_ids
89# simply generate one code span10generated_ids = model.generate(input_ids, max_length=8)11print(tokenizer.decode(generated_ids[0], skip_special_tokens=True))12# this prints "{user.username}"
Introduction
This repo provides the code for reproducing the experiments
in CodeT5: Identifier-aware Unified Pre-trained Encoder-Decoder Models for Code Understanding and Generation
. CodeT5 is a new pre-trained encoder-decoder model for programming languages, which is pre-trained on 8.35M
functions in 8 programming languages (Python, Java, JavaScript, PHP, Ruby, Go, C, and C#). In total, it achieves
state-of-the-art results on 14 sub-tasks in a code intelligence benchmark - CodeXGLUE.
The code currently includes two pre-trained checkpoints (CodeT5-small
and CodeT5-base) and scripts to fine-tune them on 4 generation tasks (
code summarization, code generation, translation, and refinement) plus 2 understanding tasks (code defect detection and
clone detection) in CodeXGLUE. We also provide their fine-tuned checkpoints to facilitate the easy replication
of our paper.
In practice, CodeT5 can be deployed as an AI-powered coding assistant to boost the productivity of software developers.
At Salesforce, we build an AI coding assistant demo using
CodeT5 as a VS Code plugin to provide three capabilities for Apex developers:
Text-to-code generation: generate code based on the natural language description.
Code autocompletion: complete the whole function of code given the target function name.
Code summarization: generate the summary of a function in natural language description.
If you find this code to be useful for your research, please consider citing:
@inproceedings{
wang2021codet5,
title={CodeT5: Identifier-aware Unified Pre-trained Encoder-Decoder Models for Code Understanding and Generation},
author={Yue Wang, Weishi Wang, Shafiq Joty, Steven C.H. Hoi},
booktitle={Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, EMNLP 2021},
year={2021},
}
@article{coderl2022,
title={CodeRL: Mastering Code Generation through Pretrained Models and Deep Reinforcement Learning},
author={Le, Hung and Wang, Yue and Gotmare, Akhilesh Deepak and Savarese, Silvio and Hoi, Steven C. H.},
journal={arXiv preprint arXiv:2207.01780},
year={2022}
}
License
The code is released under the BSD-3 License (see LICENSE.txt for details), but we also ask that users respect the
following:
This software should not be used to promote or profit from:
violence, hate, and division,
environmental destruction,
abuse of human rights, or
the destruction of people's physical and mental health.
We encourage users of this software to tell us about the applications in which they are putting it to use by emailing
codeT5@salesforce.com, and to
use appropriatedocumentation when
developing high-stakes applications of this model.
Go to sh folder, set the WORKDIR in exp_with_args.sh to be your cloned CodeT5 repository path.
You can use run_exp.py to run a broad set of experiments by simply passing the model_tag, task, and sub_task
arguments. In total, we support five models (i.e., ['roberta', 'codebert', 'bart_base', 'codet5_small', 'codet5_base'])
and six tasks (i.e., ['summarize', 'concode', 'translate', 'refine', 'defect', 'clone']). For each task, we use
the sub_task to specify which specific datasets to fine-tne on. Below is the full list:
--task
--sub_task
Description
summarize
ruby/javascript/go/python/java/php
code summarization task on CodeSearchNet data with six PLs
model_dir: where to save fine-tuning checkpoints
res_dir: where to save the performance results
summary_dir: where to save the training curves
data_num: how many data instances to use, the default -1 is for using the full data
gpu: the index of the GPU to use in the cluster
You can also revise the suggested
arguments here or directly customize the exp_with_args.sh bash file.
Please refer to the argument flags in configs.py for the full
available options. The saved training curves in summary_dir can be visualized using tensorboard.
Note that we employ one A100 GPU for all fine-tuning experiments.
How to reproduce the results using the released finetuned checkpoints?
Remove the --do_train --do_eval --do_eval_bleu and reserve only --do_test at here.
Pass the path of your downloaded finetuned checkpoint to load at here, e.g., file = "CodeT5/finetuned_models/summarize_python_codet5_base.bin"
Run the program: python run_exp.py --model_tag codet5_base --task summarize --sub_task python
How to fine-tune on your own task and dataset?
If you want to fine-tune on your dataset, you can add your own task and sub_task in configs.py (here) and add your data path and the function to read in utils.py (here and here). The read function can be implemented in _utils.py similar to this one. If your task to add is a generation task, you can simply reuse or customize the run_gen.py. For understanding tasks, please refer to run_defect.py and run_clone.py.
Get Involved
Please create a GitHub issue if you have any questions, suggestions, requests or bug-reports. We welcome PRs!