AST-T5: Pré-treinamento com Consciência Estrutural para Geração e Compreensão de Código

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

January 5, 2024

Autores: Linyuan Gong, Mostafa Elhoushi, Alvin Cheung

cs.AI

Resumo

Modelos de linguagem de grande escala (LLMs) têm feito avanços significativos em tarefas relacionadas a código, mas muitos LLMs tratam o código como sequências simples, negligenciando sua natureza estruturada. Apresentamos o AST-T5, um novo paradigma de pré-treinamento que aproveita a Árvore de Sintaxe Abstrata (AST) para aprimorar a geração, transpilação e compreensão de código. Usando programação dinâmica, nossa Segmentação Consciente da AST preserva a estrutura do código, enquanto nosso objetivo de Corrupção de Span Consciente da AST capacita o modelo a reconstruir várias estruturas de código. Diferente de outros modelos, o AST-T5 evita análises de programa complexas ou mudanças arquitetônicas, integrando-se perfeitamente com qualquer Transformer de codificador-decodificador. Avaliações mostram que o AST-T5 supera consistentemente LMs de tamanho similar em várias tarefas relacionadas a código. A consciência estrutural torna o AST-T5 particularmente poderoso em tarefas de código-para-código, superando o CodeT5 em 2 pontos na pontuação de correspondência exata para a tarefa Bugs2Fix e em 3 pontos na pontuação de correspondência exata para Transpilação Java-C# no CodeXGLUE. Nosso código e modelo estão disponíveis publicamente em https://github.com/gonglinyuan/ast_t5.

English

Large language models (LLMs) have made significant advancements in code-related tasks, yet many LLMs treat code as simple sequences, neglecting its structured nature. We introduce AST-T5, a novel pretraining paradigm that leverages the Abstract Syntax Tree (AST) for enhanced code generation, transpilation, and understanding. Using dynamic programming, our AST-Aware Segmentation retains code structure, while our AST-Aware Span Corruption objective equips the model to reconstruct various code structures. Unlike other models, AST-T5 avoids intricate program analyses or architectural changes, so it integrates seamlessly with any encoder-decoder Transformer. Evaluations show that AST-T5 consistently outperforms similar-sized LMs across various code-related tasks. Structure-awareness makes AST-T5 particularly powerful in code-to-code tasks, surpassing CodeT5 by 2 points in exact match score for the Bugs2Fix task and by 3 points in exact match score for Java-C# Transpilation in CodeXGLUE. Our code and model are publicly available at https://github.com/gonglinyuan/ast_t5.