AST-T5: Pretraining Consciente de la Estructura para la Generación y Comprensión 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

Resumen

Los modelos de lenguaje de gran escala (LLMs) han logrado avances significativos en tareas relacionadas con código, aunque muchos LLMs tratan el código como simples secuencias, descuidando su naturaleza estructurada. Presentamos AST-T5, un nuevo paradigma de preentrenamiento que aprovecha el Árbol de Sintaxis Abstracta (AST) para mejorar la generación, transpilación y comprensión de código. Utilizando programación dinámica, nuestra Segmentación Consciente del AST preserva la estructura del código, mientras que nuestro objetivo de Corrupción de Segmentos Consciente del AST capacita al modelo para reconstruir diversas estructuras de código. A diferencia de otros modelos, AST-T5 evita análisis de programas complejos o cambios arquitectónicos, por lo que se integra perfectamente con cualquier Transformer de codificador-decodificador. Las evaluaciones muestran que AST-T5 supera consistentemente a modelos de lenguaje de tamaño similar en diversas tareas relacionadas con código. La conciencia estructural hace que AST-T5 sea particularmente potente en tareas de código a código, superando a CodeT5 por 2 puntos en la puntuación de coincidencia exacta para la tarea Bugs2Fix y por 3 puntos en la puntuación de coincidencia exacta para la transpilación Java-C# en CodeXGLUE. Nuestro código y modelo están disponibles públicamente en 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.