源代码大模型

我偶尔会训练并公开发布基于程序的大型神经语言模型，包括 PolyCoder。本文将介绍如何使用这些模型。

2022年10月 - PolyCoder 已在 Huggingface 上发布！

感谢 @NinedayWang，PolyCoder 现已上线 Huggingface Hub！

可用模型包括：
* NinedayWang/PolyCoder-160M
* NinedayWang/PolyCoder-0.4B
* NinedayWang/PolyCoder-2.7B

在 Huggingface 中使用时，只需运行以下命令（需要最新版本的 transformers：pip install transformers==4.23.0）：

import transformers
from transformers import AutoTokenizer, AutoModelForCausalLM

from packaging import version
assert version.parse(transformers.__version__) >= version.parse("4.23.0")

tokenizer = AutoTokenizer.from_pretrained("NinedayWang/PolyCoder-2.7B")
model = AutoModelForCausalLM.from_pretrained("NinedayWang/PolyCoder-2.7B")

例如，可以这样使用模型：

prompt = '''def binarySearch(arr, left, right, x):
    mid = (left +'''
input_ids = tokenizer.encode(prompt, return_tensors='pt')
result = model.generate(input_ids, max_length=50, num_beams=4, num_return_sequences=4)
for res in result:
    print(tokenizer.decode(res))

目录

1. 环境配置
2. 模型（包括 PolyCoder）
3. 数据集
4. 评估
5. 如何引用

环境配置

所有当前模型均使用 GPT NeoX 工具包 进行训练。首先，按照下文所述下载预训练检查点，然后通过 Docker 镜像 或我们分叉的工具包 从源码安装 来 生成代码 或 复现我们的评估。

获取检查点

PolyCoder 的训练检查点文件托管在 此公共 Zenodo 仓库。有关当前可用模型的详细信息，请参阅 此部分。模型检查点大小最高达 6GB，这也是运行它们所需的 GPU 内存量（未测试也不推荐在 CPU 上运行）。使用以下命令下载并解压检查点文件（此处以训练了 150K 步的 27 亿参数模型为例）到名为 checkpoints/ 的目录中：

mkdir checkpoints
cd checkpoints
wget https://zenodo.org/record/6363556/files/2-7B-150K.tar
tar -xvf 2-7B-150K.tar

从源码安装

我们在 此处 维护了一个 NeoX 仓库的公共分叉，其中包含我们对代码库所做的（少量）修改，以允许在分词化中处理制表符和换行符，并且包含了运行困惑度和 HumanEval 任务的说明。请注意，此仓库使用了 LM 评估工具包的分叉版本，其中包含来自 我们工作 的代码基准测试。

构建此仓库的过程应与 GPT-NeoX 几乎完全相同。你也可以使用下面提到的 Docker 镜像，但需要将最新版本的分叉仓库挂载到容器内的 /gpt-neox 目录上。设置完成后，可以使用 generate.py 入口点（如下文 所述）进行自由形式的代码生成，或者使用 此处 的命令来计算困惑度和 HumanEval 结果，如 论文 所述。

通过 Docker

一个包含 gpt-neox 仓库 略微修改版本的 基础 Docker 镜像 可通过 DockerHub 获取：

docker pull vhellendoorn/code-lms-neox:base

此镜像可与托管在 此公共 Zenodo 仓库 的检查点文件一起使用。基础 Docker 镜像大小为 5.4GB。获取检查点后，使用以下命令启动容器（如果需要，可替换为其他 GPU 设备索引）：

nvidia-docker run --rm -it -e NVIDIA_VISIBLE_DEVICES=0 --shm-size=1g --ulimit memlock=-1 --mount type=bind,src=$PWD/checkpoints,dst=/gpt-neox/checkpoints vhellendoorn/code-lms-neox:base

代码生成

以下命令可用于根据提示生成代码：

sudo ./deepy.py generate.py configs/text_generation.yml checkpoints/configs/local_setup.yml checkpoints/configs/2-7B.yml

注意： 如果使用的不是 27 亿参数模型，请将最后的配置文件替换为相应模型大小的配置（例如，small = 1.6 亿参数，medium = 4.05 亿）。

加载检查点后，你可以输入一个示例，例如 def return1():\n """Returns 1."""\n（注意空白字符），然后观察它预测 return 1（之后可能还会预测一堆其他的 returnX 方法，具体取决于采样）。

上述对 gpt-neox 的修改主要围绕允许在提示输入中使用制表符和换行符的需求。对于 交互式 模式，可以使用转义版本（\t、\n）添加这些字符；当使用基于文件的输入时，项目将读取整个文件，而不是将每一行视为一个提示。默认情况下，以下命令将创建一个交互式提示，并返回相对较短的输出（256 个令牌），采样温度为 0.5；此行为可以在 /gpt-neox/checkpoints/configs/text_generation.yml 中更改。

较低的温度（例如 0.2）将产生更一致且（对模型而言）更合理的预测；较高的温度（如默认值）可能有助于生成和评估多个候选（参见 我们的论文 获取建议）。对于后一种设置，请考虑切换到 input-file 模式，并在相应文件中提供整个代码片段（无需转义空白字符）。

多语言模型

我们训练了多个模型，基于一个涵盖 12 种编程语言的 大型代码语料库。这包括一个 27 亿参数的模型（昵称为 PolyCoder，训练了 100K 和 150K 步）、一个 4.05 亿参数的模型（100K 和 150K 步）以及一个 1.6 亿参数的模型（150K 步）。

可用模型

所有模型均可在 公共 Zenodo 仓库 获取，以 .tar 文件形式提供，文件名具有相当自解释性（例如，2-7B-100K 表示训练了 100K 步的 27 亿参数模型）。当前可用模型包括：

• GPT2 - 2.7B: 一个 32 层、2560 维的 Transformer 模型，使用 128 个序列（256K 个令牌）的批次大小进行训练。提供 100K 步和 150K 步的模型。
• 请注意，GPT-Neox 的 默认配置 已修改，将训练步数（以及相应的学习率衰减步数）从 320K 减少到 160K，以更好地匹配可用的训练资源。因此，此模型可能尚未达到其峰值性能。
• GPT2 - 0.4B: 一个 24 层、1024 维的 Transformer 模型，基于 medium 配置，使用每批次 256K 个令牌进行训练。
• GPT2 - 160M: 一个 12 层、768 维的 Transformer 模型，基于 small 配置，使用每批次 256K 个令牌进行训练。

训练过程

训练在 4 到 8 个 NVIDIA RTX 8000 GPU 上进行，主要遵循标准配置值，但为了性能启用了 "scaled-upper-triang-masked-softmax-fusion" 和 "bias-gelu-fusion"，并略微更改了批次大小（参见 模型详情）、数据划分（改为 98.9%、0.1%、1%）、初始损失尺度（2^16）以及打印/评估间隔。

下图显示了各模型训练过程中验证损失的变化曲线。

注意事项

训练好的模型存在一些已知的小限制：
- 该模型并非为解决编程问题而训练，因此在 HumanEval 等基准测试上可能表现不佳。像 Codex（为 Copilot 提供支持）这样的模型是在自然语言上预训练的，这可能增强了它们解释自然语言提示的能力；而此模型仅从代码注释中学习语言。
- 模型在到达当前文件的（预测）结尾后，似乎会开始生成一个随机的新文件。可能是训练数据中未正确添加文档结束标记。
- 空白字符对模型 非常重要，因为输入文件未进行预处理。例如，以下代码片段将产生较差的预测，因为在 Java 中，我们不会期望在顶层出现实例方法，正如该方法内两行的单层缩进（\t）所指示的那样：

public int getTotalWeight(List<Integer> weights) {\n\t// Sum weights in parallel.\n\treturn 

调整缩进后，模型会预测更合理的续写：

public int getTotalWeight(List<Integer> weights) {\n\t\t// Sum weights in parallel.\n\t\treturn 

Codex 模型讨论了对此进行控制以提高可用性；在未来的模型版本中可能值得这样做。

数据集

249GB 多语言语料库

这是用于训练 PolyCoder 的语料库。

数据集于 2021 年 10 月 9 日至 10 日夜间克隆。要挖掘类似的训练集，请参阅 数据。

文件路径列表可从以下地址下载：https://zenodo.org/record/6363556/files/index.zip。
文件中的每一行是文件路径及其 SHA-256 哈希值，以便于去重。也就是说，哈希值允许检查任何未来测试集中的文件是否已包含在训练集中。

数据收集和过滤过程的详细信息在 论文 和下文中有描述。最终过滤后的数据集统计信息如下：

数据收集与过滤

我在 2021 年 10 月克隆了 12 种流行编程语言中至少有 50 颗星的最受欢迎仓库（每种语言约 25K 个）。对于每个项目，提取属于该项目主要语言的每个文件，得到以下训练集（经过清洗后）。这个初始的、未过滤的数据集涵盖 631GB 和 3890 万个文件。

接下来，类似于 Codex 和 CodeParrot，过滤掉非常大（>1MB）和非常短（<100 个令牌）的文件，将数据集减少到 424GB。然后根据文件内容的哈希值进行去重，这又减少了约 30% 的文件，留下 249GB 的数据和 2410 万个文件。未应用任何分词过滤器；模型处理包括所有注释在内的整个文件。基于上述文件的随机 5% 子集构建了特定于代码的词汇表。

评估

关于复现我们的困惑度和 HumanEval 结果的详细说明，请参阅我们 NeoX 仓库的 公共分叉。这又利用了我们对 LM 评估工具包的扩展。

评估 Codex

要下载我们在论文中使用的测试集（12 种编程语言），请使用：

wget https://zenodo.org/record/6363556/files/unseen_test_sets.tar.gz
tar -xvzf unseen_test_sets.tar.gz

要使用 Codex 的 API 获取这些样本的困惑度结果，请使用：

export OPENAI_API_KEY=<你的 OPEN AI API 密钥>
python3 -u Evaluation/eval_codex_all.py --dirs Code-sampled100

其中 <你的 OPEN AI API 密钥> 是一个私有字符串，可以通过注册 OpenAI 的测试版 获得。

截至 2022 年 3 月，获取 API 密钥在前 3 个月是免费的，之后需要输入信用卡信息。然而，即使在输入信用卡信息后，使用我们的评估脚本也不会产生任何费用。

结果 - HumanEval

以下是 PolyCoder 在 HumanEval 基准测试 上的结果：

结果 - 多语言语言建模

以下是 PolyCoder 在 多语言测试集 上的困惑度结果：

与其他模型的比较可在论文的图 6 中找到：

引用

对代码大模型的系统评估

@article{xu2022systematic,
  title={A Systematic Evaluation of Large Language Models of Code},
  author={Xu, Frank F and Alon, Uri and Neubig, Graham and Hellendoorn, Vincent J},
  journal={arXiv preprint arXiv:2202.13169},
  year={2022}
}