Google Colab で StableLM を試す
「Google Colab」で「StableLM」を試したので、まとめました。
1. StableLM
「StableLM」は、「Stability AI」が開発したオープンソースの言語モデルです。
アルファ版は30億パラメータと70億パラメータのモデルが用意されており、今後150億パラメータから650億パラメータのモデルも用意される予定です。開発者は、「CC BY-SA-4.0ライセンス」に従って、商用または研究目的で、「StableLMベースモデル」を自由に使用することができます。
「StableLM」は、「The Pile」をベースに構築された新しい実験用データセットで学習します。その規模は「The Pile」の3倍で、1.5兆トークンのコンテンツがあります。このデータセットの豊富さにより、「StableLM」は、30億から70億のパラメータという小さなサイズにもかかわらず、会話やコーディングのタスクで驚くほど高い性能を発揮します。
また、「ファインチューニングした研究用のモデル」も公開されました。このモデルは、会話型エージェントのための最近の5つのオープンソースデータセット (Alpaca、GPT4All、Dolly、ShareGPT、HH) を使用する予定です。このモデルはAlpacaライセンスと同様に、非営利の「CC BY-NC-SA 4.0ライセンス」に従って、研究目的のみで公開されています。
2. Colabでの実行
Colabでの実行手順は、次のとおりです。
(1) メニュー「編集→ノートブックの設定」で、「ハードウェアアクセラレータ」で「GPU」の「プレミアム」を選択。
VRAMが16GBあれば、プレミアムでなくても動くかもしれません。
(2) パッケージのインストール。
# パッケージのインストール
!pip install transformers(3) トークナイザーとモデルの準備。
今回は、7Bのファインチューニングした研究用モデルを指定しています。
・stablelm-tuned-alpha-7b : 7Bのファインチューニングした研究用モデル
・stablelm-tuned-alpha-3b : 3Bのファインチューニングした研究用モデル
・stablelm-base-alpha-7b : 7Bのベースモデル
・stablelm-base-alpha-3b : 3Bのベースモデル
from transformers import AutoModelForCausalLM, AutoTokenizer
# トークナイザーとモデルの準備
tokenizer = AutoTokenizer.from_pretrained("StabilityAI/stablelm-tuned-alpha-7b")
model = AutoModelForCausalLM.from_pretrained("StabilityAI/stablelm-tuned-alpha-7b")
model.half().cuda()(4) トークン停止クラスの定義。
from transformers import StoppingCriteria
import torch
# トークン停止クラスの定義
class StopOnTokens(StoppingCriteria):
def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
stop_ids = [50278, 50279, 50277, 1, 0]
for stop_id in stop_ids:
if input_ids[0][-1] == stop_id:
return True
return False(5) システムプロンプトの定義。
# システムプロンプトの定義
system_prompt = """<|SYSTEM|># StableLM Tuned (Alpha version)
- StableLM is a helpful and harmless open-source AI language model developed by StabilityAI.
- StableLM is excited to be able to help the user, but will refuse to do anything that could be considered harmful to the user.
- StableLM is more than just an information source, StableLM is also able to write poetry, short stories, and make jokes.
- StableLM will refuse to participate in anything that could harm a human.
"""(6) 推論の実行。
プロンプトのInstructionチューニングの書式は、Alpacaの書式でなく、ChatGPTの書式 (<|SYSTEM|>、<|USER|>、<|ASSISTANT|>) のようです。
from transformers import StoppingCriteriaList
# プロンプトの準備
prompt = "What is the capital of Japan?"
prompt = f"{system_prompt}<|USER|>{prompt}<|ASSISTANT|>"
# 推論の実行
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
tokens = model.generate(
**inputs,
max_new_tokens=64,
temperature=0.7,
do_sample=True,
stopping_criteria=StoppingCriteriaList([StopOnTokens()])
)
print("response :", tokenizer.decode(tokens[0], skip_special_tokens=False))response : <|SYSTEM|># StableLM Tuned (Alpha version)
- StableLM is a helpful and harmless open-source AI language model developed by StabilityAI.
- StableLM is excited to be able to help the user, but will refuse to do anything that could be considered harmful to the user.
- StableLM is more than just an information source, StableLM is also able to write poetry, short stories, and make jokes.
- StableLM will refuse to participate in anything that could harm a human.
<|USER|>What is the capital of Japan?<|ASSISTANT|>The capital of Japan is Tokyo.<|USER|>
