gonglinyuan/ast_t5_base

Official repository: https://github.com/gonglinyuan/ast_t5

AST-T5

Paper: AST-T5: Structure-Aware Pretraining for Code Generation and Understanding

Authors: Linyuan Gong, Mostafa Elhoushi, Alvin Cheung

Use the AST-T5 Model

The AST-T5 model is readily available on the Huggingface Model Hub (https://huggingface.co./gonglinyuan/ast_t5_base). To use our AST-T5 model in PyTorch (Python 3.8+, PyTorch 1.12+ and transformers 4.36+ are prerequisites), refer to the code snippet below:

from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

model = AutoModelForSeq2SeqLM.from_pretrained("gonglinyuan/ast_t5_base", trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained("gonglinyuan/ast_t5_base", trust_remote_code=True)

input_text = r'''def fibonacci(n):
    """return n-th fibonacci number.
    fibonacci[0] = 0
    fibonacci[1] = 1
    """'''
inputs = tokenizer(
    [input_text + "<sen001>"],  # T5-style sentinel token for completion
    max_length=1024,
    truncation=True,
    add_special_tokens=True,
    return_tensors="pt"
).input_ids
outputs = model.generate(inputs, max_length=256, do_sample=False)

output_code = tokenizer.decode(outputs[0], skip_special_tokens=True, clean_up_tokenization_spaces=False)
output_code = output_code[len("<sen001>"):]  # Remove the sentinel token
print(input_text + output_code)

Note: The ast_t5_base model is not an instruct model. It works best with specific prompts like function signatures or comments, rather than general instructions such as "Please write a code to calculate the n-th fibonacci number".

Citation

If you find the code and models useful for your research, please cite the following paper:

@article{
    ast_t5,
    title={{AST}-{T}5: Structure-Aware Pretraining for Code Generation and Understanding},
    url={http://arxiv.org/abs/2401.03003},
    DOI={10.48550/arXiv.2401.03003},
    note={arXiv:2401.03003 [cs]},
    number={arXiv:2401.03003},
    publisher={arXiv},
    author={Gong, Linyuan and Elhoushi, Mostafa and Cheung, Alvin},
    year={2024}, month=jan
}