AST-T5 : Prétraitement conscient de la structure pour la génération et la compréhension de code

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

January 5, 2024

Auteurs: Linyuan Gong, Mostafa Elhoushi, Alvin Cheung

cs.AI

Résumé

Les grands modèles de langage (LLMs) ont réalisé des avancées significatives dans les tâches liées au code, mais de nombreux LLMs traitent le code comme de simples séquences, négligeant sa nature structurée. Nous présentons AST-T5, un nouveau paradigme de pré-entraînement qui exploite l'arbre de syntaxe abstraite (AST) pour améliorer la génération, la transpilation et la compréhension du code. Grâce à la programmation dynamique, notre segmentation consciente de l'AST préserve la structure du code, tandis que notre objectif de corruption de segments conscient de l'AST permet au modèle de reconstruire diverses structures de code. Contrairement à d'autres modèles, AST-T5 évite les analyses de programme complexes ou les modifications architecturales, ce qui lui permet de s'intégrer de manière transparente avec tout encodeur-décodeur Transformer. Les évaluations montrent qu'AST-T5 surpasse systématiquement les modèles de langage de taille similaire dans diverses tâches liées au code. La conscience de la structure rend AST-T5 particulièrement puissant dans les tâches de code-à-code, surpassant CodeT5 de 2 points en score de correspondance exacte pour la tâche Bugs2Fix et de 3 points en score de correspondance exacte pour la transpilation Java-C# dans CodeXGLUE. Notre code et notre modèle sont disponibles publiquement à l'adresse 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.