CodeGeeX: 一个多语言代码生成的开源AI助手

CodeGeeX:开启AI辅助编程新时代

在人工智能快速发展的今天,代码智能生成成为了一个备受关注的研究热点。近日,来自清华大学知识工程实验室(KEG)的研究团队推出了一个重量级的开源项目——CodeGeeX,这是一个拥有130亿参数的大规模多语言代码生成模型,为程序员们带来了全新的AI编程助手。

CodeGeeX的核心特性

CodeGeeX具有以下几个突出的特点:

多语言代码生成: CodeGeeX可以生成包括Python、C++、Java、JavaScript、Go等在内的多种主流编程语言的可执行程序。
跨语言代码翻译: 支持不同编程语言之间的代码片段翻译,只需一键操作即可将程序转换为目标语言,并保持较高的准确性。
定制化编程助手: CodeGeeX提供了免费的VS Code插件,支持代码补全、解释、总结等多种功能,大大提升了用户的编码体验。
开源跨平台: 所有代码和模型权重都公开可用于研究目的。CodeGeeX同时支持华为的昇腾(Ascend)和NVIDIA平台,可在单个Ascend 910、NVIDIA V100或A100上进行推理。

模型架构与训练过程

CodeGeeX采用了基于Transformer的大规模预训练语言模型架构。它是一个自左向右的自回归解码器,可以接收代码和自然语言作为输入,并预测下一个token的概率。具体来说,CodeGeeX包含40个Transformer层,自注意力块的隐藏大小为5,120,前馈层为20,480,总参数量达到130亿。模型支持的最大序列长度为2,048。

在训练数据方面,CodeGeeX的代码语料库主要来自两个部分:

开源代码数据集,如The Pile和CodeParrot。这些数据集包含了从GitHub上收集的星标数超过100的公共代码库。
直接从GitHub上抓取的补充数据,主要包括Python、Java和C++代码。为了保证数据质量,研究团队选择了至少有1颗星且大小小于10MB的代码库,并进行了一系列的数据清洗和筛选。

最终,CodeGeeX的训练语料包含了23种编程语言,共计1587亿个代码token。为了帮助模型区分不同的编程语言,研究团队在每个代码片段的开头添加了特定的语言前缀,如"# language: Python"。

CodeGeeX训练数据分布

图1: CodeGeeX训练数据中各编程语言的占比(左)和训练过程中损失值的变化(右)

CodeGeeX的训练在1,536个昇腾910 AI处理器上进行,采用了FP16格式的模型权重,仅在layer-norm和softmax层使用FP32以提高精度和稳定性。为了提高训练效率,研究团队采用了8路模型并行训练和192路数据并行训练的策略,同时启用了ZeRO-2优化器。微批量大小为16,全局批量大小达到3,072。此外,还采用了诸如元素级算子融合、快速gelu激活、矩阵乘法维度优化等技术来进一步提升训练效率。

整个训练过程持续了近两个月,从2022年4月18日持续到6月22日,期间共处理了8500亿个token,相当于完成了5轮以上的训练。

HumanEval-X:多语言程序综合基准测试

为了更好地评估代码生成模型的多语言能力,研究团队提出了一个新的基准测试——HumanEval-X。与以往仅关注语义相似性的评估方法不同,HumanEval-X评估生成程序的功能正确性。该基准测试包含820个高质量的人工编写的数据样本,涵盖Python、C++、Java、JavaScript和Go五种编程语言,每个样本都配有测试用例。

HumanEval-X支持两项主要任务:

代码生成: 模型使用函数声明和文档字符串作为输入,生成函数实现。
代码翻译: 模型使用源语言和目标语言的函数声明以及源语言的实现作为输入,生成目标语言的实现。

评估指标采用了Codex论文中提出的无偏pass@k指标,其中n=200,k∈{1,10,100}。

HumanEval-X任务示例

图2: HumanEval-X支持的任务示例。红色、绿色和蓝色分别标记了声明、文档字符串和解决方案。

CodeGeeX的性能表现

在多语言代码生成任务中,研究团队将CodeGeeX与其他开源代码生成模型进行了比较,包括来自Meta的InCoder和来自Salesforce的CodeGen。结果显示,CodeGeeX在各种编程语言上的平均性能显著优于规模较小的模型(提高7.5%~16.3%),并与规模更大的CodeGen-Multi-16B相当(平均性能54.76% vs 54.39%)。

代码生成性能比较

图3: CodeGeeX在HumanEval-X代码生成任务上的表现。左图显示了五种语言的详细pass@k (k=1,10,100)性能,右图显示了各模型在所有语言上的平均性能。

在跨语言代码翻译任务中,研究团队评估了CodeGeeX的零样本性能,以及使用XLCoST数据集微调后的CodeGeeX-13B-FT模型的性能。结果表明,不同模型在各种语言之间的翻译性能存在差异,这可能与训练语料中的语言分布有关。例如,CodeGeeX在将其他语言翻译成Python和C++方面表现出色,而CodeGen-Multi-16B则更擅长翻译成JavaScript和Go。

CodeGeeX的应用场景

作为一个强大的AI编程助手,CodeGeeX已经在多个平台上提供了插件或扩展,方便开发者在日常工作中使用:

VS Code插件: 开发者可以在VS Code的扩展市场中搜索"CodeGeeX"安装使用。该插件支持代码补全、解释、总结等功能,大大提升了编码效率。
JetBrains IDE插件: CodeGeeX支持IntelliJ IDEA、PyCharm、GoLand、CLion等JetBrains全家桶IDE,开发者可以在插件市场搜索"CodeGeeX"安装。
在线演示: 研究团队还提供了代码生成和代码翻译的在线演示,让用户可以快速体验CodeGeeX的强大功能。

此外,CodeGeeX还展示了强大的少样本学习能力。通过提供少量示例,用户可以启发模型生成所需的程序,而无需进行昂贵的微调。这一特性为CodeGeeX打造定制化编程助手提供了可能性,用户可以根据自己的需求和编码风格来调教AI助手。

未来展望与挑战

尽管CodeGeeX在多语言代码生成和翻译方面展现出了令人印象深刻的性能,但研究团队也指出了一些值得进一步探索的方向:

模型容量与多语言能力的关系: 如何帮助模型提取编程的本质共性知识,而不是简单地增加参数量来存储多种语言的知识,这是一个值得深入研究的问题。
推理能力的一致性: 虽然CodeGeeX能够解决不同语言的问题,但其在不同语言间的表现存在较大差异。如何提升模型在各种语言中的一致性推理能力,仍然是一个挑战。
少样本能力的进一步探索: CodeGeeX展示的少样本学习能力为定制化AI编程助手提供了可能。如何更好地利用这一特性,使CodeGeeX成为真正智能的编程伙伴,还需要更多的研究和实践。

结语

CodeGeeX的发布无疑为AI辅助编程领域带来了新的活力。作为一个开源项目,它不仅为研究人员提供了宝贵的资源,也为广大开发者带来了实用的工具。随着技术的不断进步和社区的共同努力,我们有理由相信,像CodeGeeX这样的AI编程助手将在未来发挥越来越重要的作用,推动软件开发效率和质量的全面提升。

对于有兴趣深入了解或使用CodeGeeX的读者,可以访问项目的GitHub仓库获取更多信息。无论你是研究人员、学生还是专业开发者,CodeGeeX都将为你的编程之旅带来新的可能性。让我们共同期待AI与人类协作编程的美好未来!