源代码大型模型

我偶尔会训练并公开发布用于程序的大型神经语言模型，其中包括 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，这也是它们运行所需的显存容量（既未测试过也不建议在 CPU 上运行）。请下载并解压一个检查点文件（此处以训练了 15 万步的 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-evaluation-harness 的一个分支版本，其中包含来自 我们的工作 的代码基准测试。

构建此仓库的过程应与 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）；而在使用文件输入时，程序会读取整个文件，而不是将每行视为一个单独的提示。默认情况下，以下命令会创建一个交互式提示，并以 0.5 的采样温度返回相对较短的输出（256 个标记）；这一行为可在 /gpt-neox/checkpoints/configs/text_generation.yml 中更改。

较低的温度（如 0.2）会产生更一致且对模型而言更合理的预测；而较高的温度（如默认值）则可能有助于生成和评估大量候选结果（有关建议，请参阅 我们的论文）。对于后者，可以考虑切换到“输入文件”模式，并在相应的文件中提供完整的代码片段（无需转义空白字符）。

多语言模型

多个模型是在涵盖 12 种编程语言的 大型语料库 上训练而成的。其中包括一个 27 亿参数的模型（昵称为 PolyCoder，分别训练了 10 万和 15 万步）、一个 4.05 亿参数的模型（10 万和 15 万步）以及一个 1.6 亿参数的模型（15 万步）。

可用模型

所有模型均可在 公共 Zenodo 仓库 中找到，以 .tar 文件的形式提供，文件名较为直观易懂（例如，2-7B-100K 表示一个参数量为 27 亿、训练步数为 10 万的模型）。目前可用的模型包括：

• GPT2 - 2.7B： 一个 32 层、隐藏层维度为 2,560 的 Transformer 模型，采用 128 条序列（共 25.6 万 token）的批量大小进行训练。提供了训练步数分别为 10 万和 15 万的版本。
  • 需要注意的是，针对该模型，GPT-Neox 的 默认配置 被修改为将训练步数（以及学习率衰减步数）从 32 万减少至 16 万，以更好地匹配现有的训练资源。因此，该模型可能尚未达到其最佳性能。
• GPT2 - 0.4B： 一个 24 层、隐藏层维度为 1,024 的 Transformer 模型，基于 medium 配置，每批次使用 25.6 万 token 进行训练。
• GPT2 - 160M： 一个 12 层、隐藏层维度为 768 的 Transformer 模型，基于 small 配置，每批次使用 25.6 万 token 进行训练。

训练过程

训练在 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 中我们通常不会期望在顶级作用域中出现实例方法，而这两行代码却显示为同一层级的缩进：

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 日夜间完成克隆。若需挖掘类似的训练集，请参阅 Data。

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

数据收集与筛选过程在 论文 中有详细描述，最终筛选后的数据集统计如下：

语言	仓库数量	大小（GB）	文件数量
C	10,749	55G	3,037,112
C#	9,511	21G	2,514,494
C++	13,726	52G	4,289,506
Go	12,371	15G	1,416,789
Java	15,044	41G	5,120,129
JavaScript	25,144	22G	1,774,174
PHP	9,960	13G	1,714,058
Python	25,446	16G	1,550,208
Ruby	5,826	4.1G	674,343
Rust	4,991	3.5G	304,842
Scala	1,497	1.8G	245,100
TypeScript	12,830	9.2G	1,441,926

数据收集与筛选

2021 年 10 月，我从 GitHub 上克隆了 12 种流行编程语言中至少拥有 50 颗星的最受欢迎仓库（每种语言约克隆 2.5 万个），并提取了每个项目中属于其主流语言的所有文件，从而得到了下面的训练集（经过清理之后）。这一初始未过滤的数据集总大小为 631GB，包含 3890 万个文件。

随后，类似于 Codex 和 CodeParrot 的做法，我们过滤掉了超大文件（超过 1MB）和超短文件（少于 100 个 token），使数据集规模缩减至 424GB。接着，我们根据文件内容的哈希值对文件进行去重，进一步减少了约 30% 的文件数量，最终得到 249GB 的数据和 2410 万个文件。我们没有应用任何分词过滤器；模型直接处理整个文件，包括所有注释。此外，我们还在上述文件的随机 5% 子集中构建了一个专用于代码的词汇表。

评估

有关如何复现我们的困惑度和 HumanEval 结果的详细说明，请参阅 NeoX 仓库的 我们的公开分支。该分支又利用了我们对 LM Evaluation Harness 的扩展版 我们的扩展。

评估 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=<YOUR OPEN AI API KEY>
python3 -u Evaluation/eval_codex_all.py --dirs Code-sampled100

其中 <YOUR OPEN AI API KEY> 是一个私密字符串，可通过注册 OpenAI 的测试版 获取。

截至 2022 年 3 月，获取 API 密钥前 3 个月是免费的，之后则需要绑定信用卡。不过，即使绑定了信用卡，使用我们的评估脚本也不会产生任何费用。

结果 - HumanEval

以下是 PolyCoder 在 HumanEval 基准测试上的表现：

模型	通过率@1	通过率@10	通过率@100
PolyCoder (160M)	2.13%	3.35%	4.88%
PolyCoder (400M)	2.96%	5.29%	11.59%
PolyCoder (2.7B)	5.59%	9.87%	17.68%
CodeParrot (110M)	3.80%	6.57%	12.78%
CodeParrot (1.5B)	3.58%	8.03%	14.96%
GPT-Neo (125M)	0.75%	1.88%	2.97%
GPT-Neo (1.3B)	4.79%	7.47%	16.30%
GPT-Neo (2.7B)	6.41%	11.27%	21.37%
GPT-J (6B)	11.62%	15.74%	27.74%
Codex (300M)	13.17%	20.37%	36.27%
Codex (2.5B)	21.36%	35.42%	59.50%
Codex (12B)	28.81%	46.81%	72.31%

结果 - 多语言语言建模

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

语言	困惑度
C	2.3464
C#	2.5832
C++	2.9189
Go	2.567
Java	2.9194
JavaScript	3.0611
PHP	3.6954
Python	3.1767
Ruby	3.9742
Rust	3.2449
Scala	3.8735
TypeScript	3.6143

与其他模型的对比见论文中的图 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}
}

Code-LMs (PolyCoder) 快速上手指南

Code-LMs 是一系列在源代码上训练的大型神经语言模型，其中最著名的是 PolyCoder。该模型支持 12 种编程语言，适用于代码生成任务。本指南将帮助您快速部署并使用 Hugging Face 版本或本地源码版本。

环境准备

系统要求

• 操作系统: Linux (推荐) 或 macOS
• GPU: 推荐使用 NVIDIA GPU。
  • 运行 2.7B 模型至少需要 6GB 显存。
  • 运行 0.4B 模型约需 2-3GB 显存。
  • 注意：官方未测试也不推荐在 CPU 上运行。
• 内存: 建议系统内存大于 16GB。

前置依赖

• Python 3.8+
• Git
• Docker (可选，用于容器化部署)
• NVIDIA Driver 和 CUDA Toolkit (若使用 GPU)

安装步骤

本项目提供两种主要使用方式：通过 Hugging Face Transformers（推荐，最简单）或通过 源码/Docker（适合复现评估或自定义生成）。

方案一：通过 Hugging Face 使用（推荐）

这是最快捷的方式，无需下载大型检查点文件，直接利用 transformers 库加载。

1. 安装依赖 确保安装指定版本的 transformers 库：

   pip install transformers==4.23.0 packaging torch
   

   (国内用户可使用清华源加速：pip install -i https://pypi.tuna.tsinghua.edu.cn/simple transformers==4.23.0 packaging torch)

2. 验证安装 无需额外配置，直接在 Python 中导入即可。可用模型包括：

   • NinedayWang/PolyCoder-160M
   • NinedayWang/PolyCoder-0.4B
   • NinedayWang/PolyCoder-2.7B

方案二：通过 Docker 或源码使用（高级）

如果您需要复现论文中的评估结果或使用特定的生成配置，请使用此方案。

选项 A：使用 Docker (推荐)

1. 拉取镜像

   docker pull vhellendoorn/code-lms-neox:base
   

   (国内用户若无法连接 DockerHub，请配置阿里云或腾讯云镜像加速器)

2. 下载模型检查点 从 Zenodo 下载所需模型（以 2.7B 为例）：

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

   (注：若下载速度慢，可尝试使用多线程下载工具如 aria2c)

3. 启动容器 将本地检查点目录挂载到容器中：

   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
   

选项 B：从源码构建

克隆修改后的 GPT-NeoX 仓库：

git clone https://github.com/frankxu2004/gpt-neox.git
cd gpt-neox
# 按照仓库内的 README 安装依赖并配置环境

基本使用

1. 使用 Hugging Face 进行代码生成

以下示例展示如何使用 PolyCoder-2.7B 模型根据提示词生成代码。

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")

# 将模型移至 GPU (如果可用)
if torch.cuda.is_available():
    model.to("cuda")

# 定义提示词 (注意保持缩进格式)
prompt = '''def binarySearch(arr, left, right, x):
    mid = (left +'''

# 编码输入
input_ids = tokenizer.encode(prompt, return_tensors='pt')
if torch.cuda.is_available():
    input_ids = input_ids.to("cuda")

# 生成代码
# num_beams=4 表示束搜索宽度，num_return_sequences=4 表示返回 4 个候选结果
result = model.generate(input_ids, max_length=50, num_beams=4, num_return_sequences=4)

# 解码并打印结果
for res in result:
    print(tokenizer.decode(res))

2. 使用命令行生成 (Docker/源码模式)

如果您使用的是 Docker 或源码部署，可以使用提供的脚本进行交互式生成。

命令示例：

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

使用说明：

• 交互模式：启动后，您可以输入包含转义字符的提示词（例如使用 \t 代表制表符，\n 代表换行符）。
• 文件模式：修改配置文件 /gpt-neox/checkpoints/configs/text_generation.yml 可切换到文件输入模式，直接读取完整代码片段。
• 温度参数 (Temperature)：
  • 默认值较高，适合生成多样化的候选代码用于评估。
  • 若需更稳定、合理的预测，建议在配置中将 temperature 设为 0.2。

重要提示

• 缩进敏感：PolyCoder 对空白字符（空格、制表符）非常敏感。请确保输入提示词的缩进符合目标语言的语法规范，否则可能导致生成质量下降。
• 模型局限：该模型仅在代码及其注释上训练，未专门针对自然语言指令微调。对于复杂的自然语言需求描述，其表现可能不如 Codex 等模型。