CodeParrot 🦜 是一个 GPT-2 模型（1.1 亿参数），专门训练用于生成 Python 代码。
from openmind import AutoTokenizer, AutoModelForCausalLM, is_torch_npu_available
from openmind_hub import snapshot_download
import torch.nn.functional as F
from torch import Tensor
import argparse
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument(
"--model_name_or_path",
type=str,
help="Path to model",
default="jeffding/codeparrot-small-openmind",
)
args = parser.parse_args()
return args
def main():
args = parse_args()
model_path = args.model_name_or_path
if is_torch_npu_available():
device = "npu:0"
else:
device = "cpu"
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForCausalLM.from_pretrained(model_path)
model = model.to(device)
prompt = "def hello_world():"
inputs = tokenizer(prompt, return_tensors="pt", return_token_type_ids=False).to(device)
out = model.generate(**inputs, max_new_tokens=80).ravel()
out = tokenizer.decode(out)
print(out)
if __name__ == "__main__":
main()您可以直接在 transformers 中加载 CodeParrot 模型和分词器：
from transformers import AutoTokenizer, AutoModelWithLMHead
tokenizer = AutoTokenizer.from_pretrained("codeparrot/codeparrot-small")
model = AutoModelWithLMHead.from_pretrained("codeparrot/codeparrot-small")
inputs = tokenizer("def hello_world():", return_tensors="pt")
outputs = model(**inputs)或使用 pipeline：
from transformers import pipeline
pipe = pipeline("text-generation", model="codeparrot/codeparrot-small")
outputs = pipe("def hello_world():")该模型在经过清洗的 CodeParrot 🦜 数据集 上进行训练，训练设置如下：
| 配置 | 值 |
|---|---|
| 批处理大小 | 192 |
| 上下文长度 | 1024 |
| 训练步数 | 150,000 |
| 梯度累积 | 1 |
| 梯度检查点 | False |
| 学习率 | 5e-4 |
| 权重衰减 | 0.1 |
| 预热步数 | 2000 |
| 学习率调度 | Cosine |
训练在 16 块 A100 (40GB) GPU 上进行。此设置下，模型大约处理了 290 亿个 tokens。
我们在 OpenAI 的 HumanEval 基准测试集上对模型进行了评估，该基准包含一系列编程挑战：
| 指标 | 值 |
|---|---|
| pass@1 | 3.80% |
| pass@10 | 6.57% |
| pass@100 | 12.78% |
pass@k 指标 表示在 k 次生成结果中至少有一个通过测试的概率。