继Stable Diffusion爆火之后,StabilityAI近期又放大招,推出了号称是革命性的编程大模型StableCode。StableCode是其首款用于编码的LLM生成式AI产品,该产品旨在帮助程序员完成日常工作。目前已发布的版本为StableCode-Completion-Alpha-3B,是一个包含30亿个参数的代码补全模型,针对多种编程语言进行了预训练,这些编程语言是基于2023年stackoverflow开发者调查的最常用语言。
from transformers import AutoModelForCausalLM, AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("stabilityai/stablecode-completion-alpha-3b")
model = AutoModelForCausalLM.from_pretrained(
"stabilityai/stablecode-completion-alpha-3b",
trust_remote_code=True,
torch_dtype="auto",
)
model.cuda()
inputs = tokenizer("import torch\nimport torch.nn as nn", return_tensors="pt").to("cuda")
tokens = model.generate(
**inputs,
max_new_tokens=48,
temperature=0.2, do_sample=True,
)
print(tokenizer.decode(tokens[0], skip_special_tokens=True))
4K上下文
from transformers import AutoModelForCausalLM, AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("stabilityai/stablecode-completion-alpha-3b-4k")
model = AutoModelForCausalLM.from_pretrained(
"stabilityai/stablecode-completion-alpha-3b-4k",
trust_remote_code=True,
torch_dtype="auto",)
model.cuda()
inputs = tokenizer("import torch\nimport torch.nn as nn", return_tensors="pt").to("cuda")
tokens = model.generate(
**inputs, max_new_tokens=48,
temperature=0.2, do_sample=True,
)
print(tokenizer.decode(tokens[0], skip_special_tokens=True))
下面是一个StableCode利用Pytorch深度学习库完成一个相对复杂的Python文件展示(灰色文本显示了StableCode的预测)。