在这篇博文中，我们将看看构建GitHub CoPilot背后的技术所需的内容，GitHub CoPilot是一个为程序员提供编码建议的应用程序。在这个逐步指南中，我们将学习如何从头开始训练一个名为CodeParrot 🦜的大型GPT-2模型。CodeParrot可以自动完成你的Python代码-在这里试试看。让我们从头开始构建它！
创建一个大型源代码数据集
我们首先需要一个大型的训练数据集。为了训练一个Python代码生成模型，我们访问了Google的BigQuery上可用的GitHub存储，并过滤出所有的Python文件。结果是一个包含2000万个文件的180GB数据集（在这里可用）。在最初的训练实验中，我们发现数据集中的重复文件严重影响了模型的性能。进一步调查数据集，我们发现：
- 0.1%的唯一文件占所有文件的15%
- 1%的唯一文件占所有文件的35%
- 10%的唯一文件占所有文件的66%
你可以在这个Twitter帖子中了解更多我们的发现。我们移除了重复文件，并应用了Codex论文中发现的相同的清理启发式方法。Codex是CoPilot背后的模型，它是一个在GitHub代码上进行微调的GPT-3模型。
清理后的数据集仍然有50GB，可以在Hugging Face Hub上找到：codeparrot-clean。有了这个数据集，我们可以设置一个新的分词器并训练一个模型。
初始化分词器和模型
首先我们需要一个分词器。让我们专门训练一个用于代码的分词器，以便它可以很好地分割代码标记。我们可以使用现有的分词器（例如GPT-2），并使用train_new_from_iterator()方法直接在我们自己的数据集上进行训练。然后将其推送到Hub。请注意，为了保持代码块的紧凑，我们省略了导入、参数解析和日志记录等代码示例中的部分。但你可以在这里找到包括预处理和下游任务评估的完整代码。
# 用于训练的迭代器
def batch_iterator(batch_size=10):
for _ in tqdm(range(0, args.n_examples, batch_size)):
yield [next(iter_dataset)["content"] for _ in range(batch_size)]
# 基础分词器
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
base_vocab = list(bytes_to_unicode().values())
# 加载数据集
dataset = load_dataset("lvwerra/codeparrot-clean", split="train", streaming=True)
iter_dataset = iter(dataset)
# 训练和保存
new_tokenizer = tokenizer.train_new_from_iterator(batch_iterator(),
vocab_size=args.vocab_size,
initial_alphabet=base_vocab)
new_tokenizer.save_pretrained(args.tokenizer_name, push_to_hub=args.push_to_hub)在Hugging Face课程中了解更多关于分词器及其构建方法。
看到那个不显眼的streaming=True参数了吗？这个小改动有很大的影响：不再下载完整的（50GB）数据集，而是在需要时流式传输单个样本，节省了大量的磁盘空间！在Hugging Face课程中了解更多有关流式传输的信息。
现在，我们初始化一个新的模型。我们将使用与GPT-2 large（15亿个参数）相同的超参数，并调整嵌入层以适应我们的新分词器，并添加一些稳定性调整。scale_attn_by_layer_idx标志确保我们按层ID缩放注意力，reorder_and_upcast_attn主要确保我们以完全精度计算注意力，以避免数值问题。我们将新初始化的模型推送到与分词器相同的仓库。
# 加载为Python代码分词化训练的codeparrot分词器
tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_name)
# 配置
config_kwargs = {"vocab_size": len(tokenizer),
"scale_attn_by_layer_idx": True,
"reorder_and_upcast_attn": True}
# 用配置加载模型并推送到hub
config = AutoConfig.from_pretrained('gpt2-large', **config_kwargs)
model = AutoModelForCausalLM.from_config(config)
model.save_pretrained(args.model_name, push_to_hub=args.push_to_hub)现在，我们已经有了一个高效的标记器和一个刚初始化的模型，我们可以开始实际的训练循环。
实现训练循环
我们使用🤗加速库进行训练，该库使我们能够将训练从笔记本电脑扩展到多GPU机器，而无需更改任何代码。我们只需要创建一个加速器并进行一些参数处理：
accelerator = 加速器()
acc_state = {str(k): str(v) for k, v in accelerator.state.__dict__.items()}
parser = HfArgumentParser(TrainingArguments)
args = parser.parse_args()
args = Namespace(**vars(args), **acc_state)
samples_per_step = accelerator.state.num_processes * args.train_batch_size
set_seed(args.seed)我们现在准备好开始训练了！让我们使用huggingface_hub客户端库克隆带有新标记器和模型的存储库。我们将为此实验创建一个新的分支。通过这样的设置，我们可以并行运行多个实验，最后将最佳结果合并到主分支中。
# 克隆模型存储库
if accelerator.is_main_process:
hf_repo = 存储库(args.save_dir, clone_from=args.model_ckpt)
# 在存储库上检出新分支
if accelerator.is_main_process:
hf_repo.git_checkout(run_name, create_branch_ok=True)我们可以直接从本地存储库加载标记器和模型。由于我们处理的是大型模型，我们可能希望在训练过程中启用梯度检查点以减少GPU内存占用。
# 加载模型和标记器
model = AutoModelForCausalLM.from_pretrained(args.save_dir)
if args.gradient_checkpointing:
model.gradient_checkpointing_enable()
tokenizer = AutoTokenizer.from_pretrained(args.save_dir)接下来是数据集。我们使用一个数据集来简化训练，该数据集产生具有固定上下文大小的示例。为了不浪费太多数据（某些样本太短或太长），我们可以使用EOS标记连接许多示例，然后对它们进行分块。
我们一起准备的序列越多，丢弃的标记比例就越小（前一个图中的灰色标记）。由于我们希望流式传输数据集而不是提前准备好所有数据，我们使用了一个IterableDataset。完整的数据集类如下所示：
class ConstantLengthDataset(IterableDataset):
def __init__(
self, tokenizer, dataset, infinite=False, seq_length=1024, num_of_sequences=1024, chars_per_token=3.6
):
self.tokenizer = tokenizer
self.concat_token_id = tokenizer.bos_token_id
self.dataset = dataset
self.seq_length = seq_length
self.input_characters = seq_length * chars_per_token * num_of_sequences
self.epoch = 0
self.infinite = infinite
def __iter__(self):
iterator = iter(self.dataset)
more_examples = True
while more_examples:
buffer, buffer_len = [], 0
while True:
if buffer_len >= self.input_characters:
break
try:
buffer.append(next(iterator)["content"])
buffer_len += len(buffer[-1])
except StopIteration:
if self.infinite:
iterator = iter(self.dataset)
self.epoch += 1
logger.info(f"Dataset epoch: {self.epoch}")
else:
more_examples = False
break
tokenized_inputs = self.tokenizer(buffer, truncation=False)["input_ids"]
all_token_ids = []
for tokenized_input in tokenized_inputs:
all_token_ids.extend(tokenized_input + [self.concat_token_id])
for i in range(0, len(all_token_ids), self.seq_length):
input_ids = all_token_ids[i : i + self.seq_length]
if len(input_ids) == self.seq_length:
yield torch.tensor(input_ids)缓冲区中的文本并行进行标记化，然后进行连接。然后，直到缓冲区为空并且过程重新开始，才会生成分块的样本。如果设置了infinite=True，数据集迭代器将从末尾重新开始。
def create_dataloaders(args):
ds_kwargs = {"streaming": True}
train_data = load_dataset(args.dataset_name_train, split="train", streaming=True)
train_data = train_data.shuffle(buffer_size=args.shuffle_buffer, seed=args.seed)
valid_data = load_dataset(args.dataset_name_valid, split="train", streaming=True)
train_dataset = ConstantLengthDataset(tokenizer, train_data, infinite=True, seq_length=args.seq_length)
valid_dataset = ConstantLengthDataset(tokenizer, valid_data, infinite=False, seq_length=args.seq_length)
train_dataloader = DataLoader(train_dataset, batch_size=args.train_batch_size)
eval_dataloader = DataLoader(valid_dataset, batch_size=args.valid_batch_size)
return train_dataloader, eval_dataloader
train_dataloader, eval_dataloader = create_dataloaders(args)在开始训练之前，我们需要设置优化器和学习率调度。我们不希望对偏差和LayerNorm权重应用权重衰减，因此我们使用一个帮助函数来排除它们。
def get_grouped_params(model, args, no_decay=["bias", "LayerNorm.weight"]):
params_with_wd, params_without_wd = [], []
for n, p in model.named_parameters():
if any(nd in n for nd in no_decay): params_without_wd.append(p)
else: params_with_wd.append(p)
return [{"params": params_with_wd, "weight_decay": args.weight_decay},
{"params": params_without_wd, "weight_decay": 0.0},]
optimizer = AdamW(get_grouped_params(model, args), lr=args.learning_rate)
lr_scheduler = get_scheduler(name=args.lr_scheduler_type, optimizer=optimizer,
num_warmup_steps=args.num_warmup_steps,
num_training_steps=args.max_train_steps,)一个重要的问题是数据和模型将如何分布在多个GPU上。这听起来很复杂，但实际上只需要一行代码就可以完成，使用🤗加速库。
model, optimizer, train_dataloader, eval_dataloader = accelerator.prepare(
model, optimizer, train_dataloader, eval_dataloader)在底层，它将使用DistributedDataParallel，这意味着一个批次被发送到每个GPU工作进程，每个工作进程都有自己的模型副本。然后计算梯度并进行聚合以更新每个工作进程上的模型。
我们还希望定期在验证集上评估模型，因此让我们编写一个函数来完成这个任务。这是以分布式方式自动完成的，我们只需要收集所有工作进程的损失。我们还希望报告困惑度。
def evaluate(args):
model.eval()
losses = []
for step, batch in enumerate(eval_dataloader):
with torch.no_grad():
outputs = model(batch, labels=batch)
loss = outputs.loss.repeat(args.valid_batch_size)
losses.append(accelerator.gather(loss))
if args.max_eval_steps > 0 and step >= args.max_eval_steps:
break
loss = torch.mean(torch.cat(losses))
try:
perplexity = torch.exp(loss)
except OverflowError:
perplexity = float("inf")
return loss.item(), perplexity.item()现在我们准备编写主要的训练循环。它看起来很像一个普通的PyTorch训练循环。在某些地方，我们使用加速库的函数而不是原生PyTorch。此外，我们在每次评估后将模型推送到分支上。
# 训练模型
model.train()
completed_steps = 0
for step, batch in enumerate(train_dataloader, start=1):
loss = model(batch, labels=batch, use_cache=False).loss
loss = loss / args.gradient_accumulation_steps
accelerator.backward(loss)
if step % args.gradient_accumulation_steps == 0:
accelerator.clip_grad_norm_(model.parameters(), 1.0)
optimizer.step()
lr_scheduler.step()
optimizer.zero_grad()
completed_steps += 1
if step % args.save_checkpoint_steps == 0:
eval_loss, perplexity = evaluate(args)
accelerator.wait_for_everyone()
unwrapped_model = accelerator.unwrap_model(model)
unwrapped_model.save_pretrained(args.save_dir, save_function=accelerator.save)
if accelerator.is_main_process:
hf_repo.push_to_hub(commit_message=f"step {step}")
model.train()
if completed_steps >= args.max_train_steps:
break当我们调用wait_for_everyone()和unwrap_model()时，我们确保所有工作进程都准备就绪，并且之前由prepare()添加的任何模型层都被删除。我们还使用了梯度累积和梯度裁剪的功能。最后，在训练完成后，我们进行最后一次评估，并保存最终模型并将其推送到中央存储库。
# 评估并保存最后一个检查点
logger.info("训练完成后进行评估和保存模型")
eval_loss, perplexity = evaluate(args)
log_metrics(step, {"loss/eval": eval_loss, "perplexity": perplexity})
accelerator.wait_for_everyone()
unwrapped_model = accelerator.unwrap_model(model)
unwrapped_model.save_pretrained(args.save_dir, save_function=accelerator.save)
if accelerator.is_main_process:
hf_repo.push_to_hub(commit_message="final model")完成！这就是使用仅150行代码从头开始训练一个完整的GPT-2模型所需的所有代码。为了使代码更加紧凑，我们没有显示脚本的导入和日志。现在让我们真正开始训练吧！
使用这段代码，我们为即将发布的《Transformer和NLP》一书训练了两个模型：一个是1.1亿参数的GPT-2模型，另一个是15亿参数的GPT-2模型。我们使用了一台16个A100 GPU的机器，分别训练了1天和1周。足够的时间喝杯咖啡和读一两本书了！
评估
对于预训练来说，这段时间还相对较短，但我们已经可以观察到与类似模型相比的良好下游性能。我们在OpenAI的HumanEval基准测试上评估了这些模型，该测试在Codex论文中介绍过。它衡量了编码模型在近200个编码挑战上的表现。
请注意，我们在大约250-300亿个标记上训练了CodeParrot，而GPT-neo是在3000亿个标记上训练的，Codex是在3000亿（GPT-3检查点）+ 1000亿（代码）个标记上训练的。因此，您在这里确实会得到很好的回报。
但是，这个pass@k指标到底是什么？简单来说，它衡量了从模型生成的k个候选解中，至少有一个程序通过了编程问题的单元测试的概率。自然而然，随着k的增加，该指标也会增加，因为更多的候选者有可能解决编码挑战。您可以在datasets中使用该指标：
from datasets import load_metric
code_eval = datasets.load_metric("code_eval")
test_cases = ["assert add(2,3)==5"]
candidates = [["def add(a,b): return a*b", "def add(a, b): return a+b"]]
pass_at_k, results = code_eval.compute(references=test_cases, predictions=candidates, k=[1, 2])
print(pass_at_k)
# -> {'pass@1': 0.5, 'pass@2': 1.0}您还可以使用datasets加载OpenAI的HumanEval数据集：
from datasets import load_dataset
dataset = load_dataset("openai_humaneval")现在让我们来看看我们在焦急等待了这么久之后得到了什么！
结果
让我们从简单的开始，检查一下我们是否可以提示模型编写一个根据路径返回文件大小的函数：
提示：
def get_files_size(filename):完成：
"""
返回文件的大小。
:param filepath: 文件的路径。
:type filepath: str
:return: 文件的大小。
:rtype: int
"""
return os.path.getsize(filepath)到目前为止，一切都还不错，但让我们再进一步尝试一些更有趣的东西。构建产品级别的代码的一个较不令人兴奋的方面之一是编写单元测试。让我们看看是否可以使用CodeParrot来为我们完成这个任务！
提示：
def is_even(value):
"""如果value是一个偶数，则返回True。"""
return value % 2 == 0
# 为is_even设置单元测试
import unittest完成：
class TestIsEven(unittest.TestCase):
"""用于测试is_even函数的测试类。"""
def test_is_even(self):
"""测试is_even函数。"""
self.assertTrue(is_even(0))
self.assertTrue(is_even(1))
self.assertTrue 尽管测试条件本身看起来不太对，但模型已经为我们构建了所有的样板代码。这一点单独就已经非常有用，可以节省我们几分钟的时间！
但我们真正想做的是构建更好的Transformer模型。让我们使用CodeParrot来帮助我们做到这一点：
提示：
from transformers import AutoTokenizer, AutoModelForSequenceClassification
# 构建一个BERT分类器完成：
tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased')
model = AutoModelForSequenceClassification.from_pretrained('bert-base-uncased')看几个例子很好，但现在你可能想自己尝试一下。我们专门为此准备了几个在线演示！
演示
你可以在 Hugging Face Spaces 上玩弄这些模型的两个演示。第一个可以快速生成模型代码，第二个可以用模型突出显示代码中的错误。
- 代码生成
- 代码突出显示
你还可以直接使用 transformers 库中的模型：
from transformers import pipeline
pipe = pipeline('text-generation', model='lvwerra/codeparrot')
pipe('def hello_world():')总结
在这篇简短的博文中，我们介绍了训练一个名为 CodeParrot 🦜 的大型 GPT-2 模型用于代码生成的所有步骤。使用 🤗 Accelerate，我们构建了一个不到 200 行代码的训练脚本，可以轻松地在多个GPU上扩展。现在你可以训练自己的 GPT-2 模型了！
本文简要介绍了 CodeParrot 🦜，但如果你对如何预训练这些模型感兴趣，我们建议阅读即将发布的关于 Transformers 和 NLP 的专门章节。该章节提供了关于构建自定义数据集、训练新的分词器和架构选择等方面的更多细节。