AST-T5: Pre-addestramento Strutturale per la Generazione e Comprensione del Codice

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

January 5, 2024

Autori: Linyuan Gong, Mostafa Elhoushi, Alvin Cheung

cs.AI

Abstract

I modelli linguistici di grandi dimensioni (LLM) hanno compiuto progressi significativi nelle attività legate al codice, tuttavia molti LLM trattano il codice come semplici sequenze, trascurandone la natura strutturata. Introduciamo AST-T5, un nuovo paradigma di pre-addestramento che sfrutta l'Albero Sintattico Astratto (AST) per migliorare la generazione, la trasposizione e la comprensione del codice. Utilizzando la programmazione dinamica, la nostra Segmentazione Consapevole dell'AST mantiene la struttura del codice, mentre l'obiettivo di Corruzione di Span Consapevole dell'AST consente al modello di ricostruire varie strutture di codice. A differenza di altri modelli, AST-T5 evita complesse analisi di programma o modifiche architetturali, integrandosi quindi perfettamente con qualsiasi Transformer encoder-decoder. Le valutazioni dimostrano che AST-T5 supera costantemente modelli linguistici di dimensioni simili in varie attività legate al codice. La consapevolezza della struttura rende AST-T5 particolarmente potente nelle attività di codice-a-codice, superando CodeT5 di 2 punti nel punteggio di corrispondenza esatta per il task Bugs2Fix e di 3 punti nel punteggio di corrispondenza esatta per la trasposizione Java-C# in CodeXGLUE. Il nostro codice e il modello sono disponibili pubblicamente all'indirizzo 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.