starcoder - 支持80多种编程语言的代码生成工具

StarCoder 项目介绍

项目简介

StarCoder 是一个强大的代码生成模型，拥有 155 亿个参数，支持 80 多种编程语言，训练数据来自于筛选后的 The Stack 数据集（v1.2）。该模型的训练采用了多查询注意力机制和一个长达 8192 词元的上下文窗口。其目标是在一个包含 1 万亿词元的数据集上进行填补并生成代码。StarCoder 项目的相关资源包括 GitHub 仓库 bigcode/Megatron-LM，以及项目网站和学术论文等。

使用方法

预期用途

StarCoder 模型主要用于 GitHub 上的代码，因此其并不是一个指令模型，对诸如“写一个平方根计算函数”这样的命令不太适用。但通过使用 Tech Assistant 提示，可以将其转变为一个技术助手。

代码生成示例

# 安装 transformers
from transformers import AutoModelForCausalLM, AutoTokenizer

checkpoint = "bigcode/starcoder"
device = "cuda" # 如使用 GPU 或 "cpu"

tokenizer = AutoTokenizer.from_pretrained(checkpoint)
model = AutoModelForCausalLM.from_pretrained(checkpoint).to(device)

inputs = tokenizer.encode("def print_hello_world():", return_tensors="pt").to(device)
outputs = model.generate(inputs)
print(tokenizer.decode(outputs[0]))

填补中间

使用特定标记来识别输入和输出中的前缀、中间和后缀部分：

input_text = "<fim_prefix>def print_hello_world():\n    <fim_suffix>\n    print('Hello world!')<fim_middle>"
inputs = tokenizer.encode(input_text, return_tensors="pt").to(device)
outputs = model.generate(inputs)
print(tokenizer.decode(outputs[0]))