AIGC代码生成模型之Salesforce/codegen-6B-mono

CodeGen是一系列用于程序合成的自回归语言模型，摘自Erik Nijkamp、Bo、Hiroaki Hayashi、Lifu Tu、Huan Wang、Yingbo Zhou、Silvio Savarese、Caiming Xiong的论文《程序合成的对话范式》。模型最初在该存储库中发布，分为3种预训练数据变体（NL、Multi-、Mono）和4种模型大小变体（350M、2B、6B、16B）。

该存储库中包含的checkpoint在本文中表示为CodeGen Mono 6B，其中“Mono”表示使用CodeGen Multi 6B初始化模型，并在Python编程语言数据集上进一步预训练，“6B”表示可训练参数的数量。

GitHub: https://github.com/salesforce/CodeGen

数据训练

这个checkpoint（CodeGen Mono 6B）首先用CodeGen Multi 6B初始化，然后在BigPython数据集上进行预训练。该数据由71.7亿个Python编程语言的令牌组成。有关更多详细信息，请参阅本文第2.1节。

训练流程

CodeGen使用交叉熵损失进行训练，以最大限度地提高顺序输入的可能性。该系列模型由谷歌使用多个TPU-v4-512进行训练，利用数据和模型并行性。有关更多详细信息，请参阅本文第2.3节。

结果评估

两个代码生成基准上评估我们的模型：HumanEval和MTPB。有关更多详细信息，请参阅论文。

使用场景和限制

作为一种自回归语言模型，CodeGen能够从给定的自然语言和编程语言文本中提取特征，并计算它们的可能性。然而，该模型旨在并最擅长程序合成，即在给定英文提示的情况下生成可执行代码，其中提示应以注释字符串的形式出现。该模型也可以完成部分生成的代码。

如何使用

使用AutoModelForCausalLM功能可以轻松加载此模型：

from transformers import AutoTokenizer, AutoModelForCausalLM
tokenizer = AutoTokenizer.from_pretrained("Salesforce/codegen-6B-mono")
model = AutoModelForCausalLM.from_pretrained("Salesforce/codegen-6B-mono")

text = "def hello_world():"
input_ids = tokenizer(text, return_tensors="pt").input_ids

generated_ids = model.generate(input_ids, max_length=128)
print(tokenizer.decode(generated_ids[0], skip_special_tokens=True))