StabilityAIのStableCodeをローカル環境で試してみる

Romot

StabiliyAIがコード補完をおこなえるLLM「StableCode」をリリースしたらしいので、
ローカルで動くかどうか試してみます。

Stability AIがオープンソースのコーディング補助AI「StableCode」を発表

環境:

GPU: RTX3090 24GB
OS: Windows10

Python 3.10系

Romot

以下のうち、StableCode-Instruct-Alpha-3Bを試してみます。

StableCode-Completion-Alpha-3B-4K
ベースモデル

StableCode-Instruct-Alpha-3B
ベースモデルをファインチューンしたもので、指示に従いコードを生成できる

StableCode-Completion-Alpha-3B
オートコンプリートに特化したモデル

Romot

gitでモデルをローカルにcloneしてみます。

cd your_dir_name
git lfs install
start-ssh-agent
git clone git@hf.co:stabilityai/stablecode-instruct-alpha-3b

Romot

hugging-faceにあるサンプルを調整し、
run_sample.pyとして以下を作成して実行します。

from transformers import AutoModelForCausalLM, AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("stabilityai/stablecode-instruct-alpha-3b")
model = AutoModelForCausalLM.from_pretrained(
  "stabilityai/stablecode-instruct-alpha-3b",
  trust_remote_code=True,
  torch_dtype="auto",
)
model.cuda()
inputs = tokenizer(
  "###Instruction\nGenerate a python function to find number of CPU cores###Response\n",
  return_tensors="pt",
  return_token_type_ids=False,
).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))

Romot

以下のようなコードが生成されました。

###Instruction
Generate a python function to find number of CPU cores

###Response
def get_cpu_count():
    """Returns the number of CPU cores"""
    # Get the number of CPU cores
    cpu_count = multiprocessing.cpu_count()
    return cpu_count

Romot

かんたんにpython run.py "PROMPT"で実行できるようにしてみます。
トークン数もサンプルでは48のようですが4096に修正しています。

import sys
from transformers import AutoModelForCausalLM, AutoTokenizer

# Ensure that a command-line argument was provided
if len(sys.argv) < 2:
    print("Please provide a code instruction as an argument.")
    sys.exit()

# Extract instruction from command-line argument
instruction = sys.argv[1]
formatted_input = f"###Instruction\n\n{instruction}\n\n###Response\n\n"

# Load tokenizer and model
tokenizer = AutoTokenizer.from_pretrained("stabilityai/stablecode-instruct-alpha-3b")
model = AutoModelForCausalLM.from_pretrained(
    "stabilityai/stablecode-instruct-alpha-3b",
    trust_remote_code=True,
    torch_dtype="auto",
)
model.cuda()

# Tokenize and generate response
inputs = tokenizer(formatted_input, 
    return_tensors="pt",
    return_token_type_ids=False,
).to("cuda")
tokens = model.generate(
    **inputs,
    max_new_tokens=4096, // 48 to 4096
    temperature=0.2,
    do_sample=True,
)

# Print generated response
print(tokenizer.decode(tokens[0], skip_special_tokens=True))

Romot

カウントアップする簡単なReactコンポーネントを書かせてみます。
とりあえず動いた…！

速度は計測していませんが、感覚的にはGPT-4による生成の2倍ぐらいでした。
RTX3090(24GB)環境ではOUT_OF_MEMORYになることもないようです。

generate React(TypeScript) component code of count up button

###Instruction

generate React(TypeScript) component code of count up button

###Response


import React, { useState } from'react';

const CountUpButton = () => {
  const [count, setCount] = useState(0);

  const handleClick = () => {
    setCount(count + 1);
  };

  return (
    <button onClick={handleClick}>
      You clicked {count} times
    </button>
  );
};

export default CountUpButton;