输入注释，tab 一下，几十行代码就出来了。

这便是我们研发的大规模多语言代码生成模型 CodeGeeX。

 

一、原理

 

模型

CodeGeeX是一个基于transformers的大规模预训练编程语言模型。它是一个从左到右生成的自回归解码器，将代码或自然语言标识符（token）作为输入，预测下一个标识符的概率分布。

CodeGeeX含有40个transformer层，每层自注意力块的隐藏层维数为5120，前馈层维数为20480，总参数量为130亿。模型支持的最大序列长度为2048，足以涵盖大部分使用场景。

 

语料

CodeGeeX的训练语料由两部分组成。

第一部分是开源代码数据集，The Pile与CodeParrot。The Pile包含GitHub上拥有超过100颗星的一部分开源仓库，我们从中选取了23种编程语言的代码。

第二部分是补充数据，直接从GitHub开源仓库中爬取Python、Java、C++代码。

为了让模型区分不同语言，我们在每个样本的开头加上一个前缀，其形式为[注释符] language: [语言]，例如：# language: Python。我们对数据进行了去重和清洗，整个代码语料含有23种编程语言、总计1587亿个标识符（不含填充符）。

平台

CodeGeeX 是由清华大学知识工程实验室研发，其中鹏城实验室提供了算力支持，智谱AI、华为MindSpore提供了技术支持。

具体而言，CodeGeeX 采用了华为 MindSpore 框架来实现，使用了鹏城实验室“鹏城云脑II”平台中 192 节点昇腾910 AI处理器，在 20 多种编程语言的代码语料库历时两个月训练而成。

此外，我们还将模型适配到了其他平台上，目前能同时支持昇腾和英伟达平台。

 

二、功能

如前所述，CodeGeeX采用的是从左到右生成的自回归解码，因此CodeGeeX 能够完成许多生成类的编程问题，例如「根据注释生成代码」、「根据代码生成解释」、「根据上下文做代码补全」、「不同语言之间的代码翻译」等等，你完全可以根据自己的需求，输入相应的提示，从而生成有意思的内容。

我们这里简要介绍两类功能：代码生成和代码翻译。

 

代码生成

代码生成，包含了「根据注释生成代码」、「根据上下文做代码补全」等生成类的能力。

目前，CodeGeeX 能够支持生成Python、C++、Java、JavaScript和Go等多种主流编程语言的代码。在HumanEval-X代码生成任务（我们会在另外一篇文章中介绍）上取得47%~60%求解率，较其他开源基线模型有更佳的平均性能。

这是在HumanEval-X中，CodeGeeX在 5 种语言下对各问题的解答率，我们按Python下的解答率做了排序：

我们观察了每道题目的具体解答情况，发现题目在不同语言上的解答率有着较大区别，某些题目更适合用特定语言进行解答。

 

跨语言代码翻译

有时候，我们会期望把一种语言实现的代码，转换成另一种语言的代码，理解和重敲，很耗费精力。

CodeGeeX 可以实现一键式翻译，和自然语言的翻译一样丝滑：

我们评测了模型在多语言间代码翻译上的性能：

上图显示了 4 种模型在 5 种语言之间的翻译性能。

• 4 种模型分别为：

  • InCoder-6.7B

  • CodeGen-Multi-16B

  • CodeGeeX-13B

  • CodeGeeX-13B-FT（即经过微调的CodeGeeX）；

• 5 种语言分别为：

  • Python

  • C++

  • Java

  • JavaScript

  • Go

•  

上表显示出，模型对特定语言存在偏好，比如CodeGeeX擅长将其他语言翻译为Python与C++，而CodeGen-Multi-16B擅长翻译为JavaScript和Go，这可能是由于训练集中的语料占比存在差异。

 

在20个翻译对中，我们还观察到两种语言互相翻译的表现常常是呈负相关的，这可能说明现有的模型还不足以学好所有的语言。

 

三、插件

基于 CodeGeeX 的强大功能，我们开发了 VS Code 插件，在 VS Code 插件市场搜索「codegeex」即可找到，该插件可以完全免费下载使用。VS Code 版本的 CodeGeeX插件，我们后续同样将完全开源出来。任何人都可以在此基础上，做出更进一步的优化。

其他 IDE 平台的插件也正在开发当中，很快便可以上线。

目前CodeGeex支持在VS Code市场和Jetbrains IDEs下载使用。

VS Code插件安装地址：https://marketplace.visualstudio.com/items?itemName=aminer.codegeex

Jetbrains IDEs插件安装地址：https://plugins.jetbrains.com/plugin/20587-codegeex

---

目前CodeGeeX，我们有以下几个开放形式：

1）代码、模型，完全开源开放：https://models.aminer.cn/codegeex

2）在线演示 demo：

代码生成：https://models.aminer.cn/codegeex/playground

代码翻译：https://models.aminer.cn/codegeex/codeTranslator

3）提供VS code插件：

https://marketplace.visualstudio.com/items?itemName=aminer.codegeex

其他平台插件，如果感兴趣，我们可以一起合作来开发

4）提供 API 接口：

https://tianqi.aminer.cn/open/document/code_ref/codegeex_generation

5）论文将于近期发布。