AST-T5: Структурно-ориентированное предобучение для генерации и понимания кода

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

January 5, 2024

Авторы: Linyuan Gong, Mostafa Elhoushi, Alvin Cheung

cs.AI

Аннотация

Крупные языковые модели (LLMs) достигли значительных успехов в задачах, связанных с кодом, однако многие из них обрабатывают код как простые последовательности, игнорируя его структурированную природу. Мы представляем AST-T5 — новую парадигму предварительного обучения, которая использует абстрактное синтаксическое дерево (AST) для улучшенной генерации, трансляции и понимания кода. Наш метод AST-Aware Segmentation, основанный на динамическом программировании, сохраняет структуру кода, а задача AST-Aware Span Corruption позволяет модели восстанавливать различные структуры кода. В отличие от других моделей, AST-T5 не требует сложного анализа программ или изменений архитектуры, что позволяет ей легко интегрироваться с любым кодировщиком-декодировщиком на основе Transformer. Оценки показывают, что AST-T5 стабильно превосходит языковые модели аналогичного размера в различных задачах, связанных с кодом. Осознание структуры делает AST-T5 особенно эффективной в задачах преобразования кода, превосходя CodeT5 на 2 балла по точному совпадению в задаче Bugs2Fix и на 3 балла в задаче трансляции Java-C# в CodeXGLUE. Наш код и модель доступны по адресу 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.