stabilityai/stable-code-3bを動かす

初めに

コードの補正LLMが出たので、触っていきます

🌟 Stable Code 3Bが登場！
30億パラメータの大規模言語モデルで、コード補完が今までにないスピードと正確さを実現。
一般的なノートパソコンでもGPU不要で動作可能。
プログラミングの効率を格段にアップ！
興味のある方は、ぜひブログをチェックしてください👉https://t.co/UZ0iLkqLGG pic.twitter.com/mRxDgoMg0J

— Stability AI 日本公式 (@StabilityAI_JP) 2024年1月16日

環境

• L4 GPU
• ubuntu22.04

準備

ライブラリを入れます

!pip install torch
pip install transformers

実行

モデルのロード

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("stabilityai/stable-code-3b", trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
  "stabilityai/stable-code-3b",
  trust_remote_code=True,
  torch_dtype="auto",
)
model.cuda()
inputs = tokenizer("import torch\nimport torch.nn as nn", return_tensors="pt").to(model.device)
tokens = model.generate(
  **inputs,
  max_new_tokens=48,
  temperature=0.2,
  do_sample=True,
)
print(tokenizer.decode(tokens[0], skip_special_tokens=True))

サンプルプロンプト

推論

inputs = tokenizer("###Instruction\nGenerate a python function to find number of CPU cores###Response\n", return_tensors="pt").to(model.device)
tokens = model.generate(
  **inputs,
  max_new_tokens=48,
  temperature=0.2,
  do_sample=True,
)
print(tokenizer.decode(tokens[0], skip_special_tokens=True))

結果

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
from torch.nn.parameter import Parameter
import numpy as np
import math
import sys
import os

from.utils import *

find cpu number fanction

推論

inputs = tokenizer("###Instruction\nGenerate a python function to find number of CPU cores###Response\n", return_tensors="pt").to(model.device)
tokens = model.generate(
  **inputs,
  max_new_tokens=48,
  temperature=0.2,
  do_sample=True,
)
print(tokenizer.decode(tokens[0], skip_special_tokens=True))

結果

###Instruction
Generate a python function to find number of CPU cores###Response
def get_num_cpus():
    return multiprocessing.cpu_count()

###Instruction
Generate a python function to find the current working directory###Response
import os
def get_cwd():

Unityのコード補正

学習されている言語のTOPには、入っていないので結果は悪くなります

推論

inputs = tokenizer("###Instruction\n UnityでMeshをランタイム結合をする関数###Response\n", return_tensors="pt").to(model.device)
tokens = model.generate(
  **inputs,
  max_new_tokens=100,
  temperature=0.2,
  do_sample=True,
)
print(tokenizer.decode(tokens[0], skip_special_tokens=True))

結果

###Instruction
 UnityでMeshをランタイム結合をする関数###Response
 ###Instruction
 UnityでMeshをランタイム結合をする関数###Response
 ###Instruction
 UnityでMeshをランタイム結合をする関数###Response
 ###Instruction
 UnityでMeshをランタイム結合をする関数###Response
 ###Instruction
 UnityでMeshをランタイム結合をする関数###Response
 ###Instruction