一种基于多任务学习的代码补全方法

作者：帅祥; 魏乐^*; 舒红平

来源：计算机应用研究, 2023, 40(03): 863-867.

DOI：10.19734/j.issn.1001-3695.2022.08.0396

• 代码补全
• 多任务学习
• 抽象语法树
• Transformer
• 集束搜索

针对基于语言模型的代码补全方法忽略源代码中结构信息和命名信息而导致补全准确率偏低的问题，提出一种基于多任务学习的代码补全方法(multi-task learning code completion, MTLCC)。MTLCC对源代码数据集进行数据清洗和预处理，通过抽象语法树(abstract syntax tree, AST)提取源代码中的结构信息和命名信息；构建基于Transformer的多任务学习网络，采用软参数共享学习方式分别对源代码的节点类型预测、节点值预测和语句预测进行训练；利用集束搜索对模型推理结果进行后期处理生成代码补全建议列表。实验结果表明，与Pointer network、LSTM+TransformerXL和CodeGPT中最好的模型相比，MTLCC在节点类型预测和节点值预测任务中的准确率分别提升了2.5%和1.7%、MRR分别提升了3.7%和2.4%,在语句预测任务中的ROUGE-L分数(L@4)提升了0.055。

• 单位
  成都信息工程大学