|
--- |
|
license: apache-2.0 |
|
tags: |
|
datasets: |
|
- code_search_net |
|
--- |
|
|
|
# CodeT5 (base-sized model) |
|
|
|
Pre-trained CodeT5 model. It was introduced in the paper [CodeT5: Identifier-aware Unified Pre-trained Encoder-Decoder Models |
|
for Code Understanding and Generation](https://arxiv.org/abs/2109.00859) by Yue Wang, Weishi Wang, Shafiq Joty, Steven C.H. Hoi and first released in [this repository](https://github.com/salesforce/CodeT5). |
|
|
|
Disclaimer: The team releasing CodeT5 did not write a model card for this model so this model card has been written by the Hugging Face team (more specifically, [nielsr](https://huggingface.co/nielsr)). |
|
|
|
## Model description |
|
|
|
From the abstract: |
|
|
|
"We present CodeT5, a unified pre-trained encoder-decoder Transformer model that better leverages the code semantics conveyed from the developer-assigned identifiers. Our model employs a unified framework to seamlessly support both code understanding and generation tasks and allows for multi-task learning. Besides, we propose a novel identifier-aware pre-training task that enables the model to distinguish which code tokens are identifiers and to recover them when they are masked. Furthermore, we propose to exploit the user-written code comments with a bimodal dual generation task for better NL-PL alignment. Comprehensive experiments show that CodeT5 significantly outperforms prior methods on understanding tasks such as code defect detection and clone detection, and generation tasks across various directions including PL-NL, NL-PL, and PL-PL. Further analysis reveals that our model can better capture semantic information from code." |
|
|
|
## Intended uses & limitations |
|
|
|
You can use the model to fine-tune it on code understanding tasks, such as . See the [model hub](https://huggingface.co/models?search=google/vit) to look for |
|
fine-tuned versions on a task that interests you. |
|
|
|
### How to use |
|
|
|
Here is how to use this model: |
|
|
|
```python |
|
from transformers import RobertaTokenizer, T5ForConditionalGeneration |
|
|
|
tokenizer = RobertaTokenizer.from_pretrained('Salesforce/codet5-base') |
|
model = T5ForConditionalGeneration.from_pretrained('Salesforce/codet5-base') |
|
|
|
text = "def greet(user): print(f'hello <extra_id_0>!') </s>" |
|
inputs = tokenizer(text, return_tensors="pt").input_ids |
|
|
|
# simply generate a single sequence |
|
generated_ids = model.generate(input_ids, max_length=8) |
|
print(tokenizer.decode(generated_ids[0], skip_special_tokens=True)) |
|
# this prints {user.name} |
|
|
|
# or, generating 20 sequences with maximum length set to 10 |
|
outputs = model.generate(input_ids=input_ids, |
|
num_beams=200, num_return_sequences=20, |
|
max_length=10) |
|
|
|
_0_index = text.index('<extra_id_0>') |
|
_result_prefix = text[:_0_index] |
|
_result_suffix = text[_0_index+12:] # 12 is the length of <extra_id_0> |
|
|
|
def _filter(output, end_token='<extra_id_1>'): |
|
# The first token is <pad> (indexed at 0), the second token is <s> (indexed at 1) |
|
# and the third token is <extra_id_0> (indexed at 32099) |
|
# So we only decode from the fourth generated id |
|
_txt = tokenizer.decode(output[3:], skip_special_tokens=False, clean_up_tokenization_spaces=False) |
|
if end_token in _txt: |
|
_end_token_index = _txt.index(end_token) |
|
return _result_prefix + _txt[:_end_token_index] + _result_suffix |
|
else: |
|
return _result_prefix + _txt + _result_suffix |
|
|
|
results = list(map(_filter, outputs)) |
|
print(results) |
|
# this prints: |
|
#["def greet(user): print(f'hello {user.name} {user!') </s>", |
|
# "def greet(user): print(f'hello {user.username} {user!') </s>", |
|
# "def greet(user): print(f'hello {user.name}: {user!') </s>", |
|
# "def greet(user): print(f'hello {user}') print(f!') </s>", |
|
# "def greet(user): print(f'hello {user.name} �!') </s>", |
|
# "def greet(user): print(f'hello {user}') print ( f!') </s>", |
|
# "def greet(user): print(f'hello {user.username}: {user!') </s>", |
|
# "def greet(user): print(f'hello {user}' ) print(f!') </s>", |
|
# "def greet(user): print(f'hello {user.username} �!') </s>", |
|
# "def greet(user): print(f'hello {user.name}, {user!') </s>", |
|
# "def greet(user): print(f'hello {user.login} {user!') </s>", |
|
# "def greet(user): print(f'hello {user} →!') </s>", |
|
# "def greet(user): print(f'hello {user}!') print(!') </s>", |
|
# "def greet(user): print(f'hello {user.name} ({user!') </s>", |
|
# "def greet(user): print(f'hello {user.email} {user!') </s>", |
|
# "def greet(user): print(f'hello {user}!') print (!') </s>", |
|
# "def greet(user): print(f'hello {user.username}, {user!') </s>", |
|
# "def greet(user): print(f'hello {user}' ) print ( f!') </s>", |
|
# "def greet(user): print(f'hello {user.nickname} {!') </s>", |
|
# "def greet(user): print(f'hello {user} {user.name!') </s>"] |
|
``` |
|
|
|
## Training data |
|
|
|
The CodeT5 model was pretrained on CodeSearchNet [Husain et al., 2019](https://arxiv.org/abs/1909.09436). Additionally, the authors collected two datasets of C/CSharp from [BigQuery1](https://console.cloud.google.com/marketplace/details/github/github-repos) to ensure that all downstream tasks have overlapped programming languages with the pre-training data. In total, around 8.35 million instances are used for pretraining. |
|
|
|
## Training procedure |
|
|
|
### Preprocessing |
|
|
|
This model uses a code-specific BPE (Byte-Pair Encoding) tokenizer. One can prepare text (or code) for the model using RobertaTokenizer, with the files from this repository. |
|
|
|
## Evaluation results |
|
|
|
For evaluation results on several downstream benchmarks, we refer to the paper. |
|
|
|
### BibTeX entry and citation info |
|
|
|
```bibtex |
|
@misc{wang2021codet5, |
|
title={CodeT5: Identifier-aware Unified Pre-trained Encoder-Decoder Models for Code Understanding and Generation}, |
|
author={Yue Wang and Weishi Wang and Shafiq Joty and Steven C. H. Hoi}, |
|
year={2021}, |
|
eprint={2109.00859}, |
|
archivePrefix={arXiv}, |
|
primaryClass={cs.CL} |
|
} |
|
``` |
