【お手軽】サイバーエージェントの日本語特化型の大規模言語モデルOpenCALMをGoogleColaboratoryで再現

はじめに

この記事はCyberAgentが公開した日本語特化のLLM(大規模言語モデル)、OpenCALMをとりあえず試したメモです。

結論

以下のコードで実行できます。

準備

# パッケージのインストール
!pip install transformers accelerate

用意

# パッケージのインポート
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# モデルとトークナイザーの準備
model = AutoModelForCausalLM.from_pretrained(
  "cyberagent/open-calm-3b", 
   device_map="auto", 
   torch_dtype=torch.float16
)
tokenizer = AutoTokenizer.from_pretrained("cyberagent/open-calm-3b")

実行方法

# プロンプトの準備
prompt = "Q:睡眠で大事なのは、量、タイミング、質のどれ？\nA:"

# 推論の実行
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
with torch.no_grad():
    tokens = model.generate(
        **inputs,
        max_new_tokens=64,
        do_sample=True,
        temperature=0.7,
        pad_token_id=tokenizer.pad_token_id,
    )
    
output = tokenizer.decode(tokens[0], skip_special_tokens=True)
print(output)

結果

メモ

• 事前に「編集」⇒「ノートブックの設定」でGPUにする必要がある。
  

• cyberagent/open-calm-7bは無料のT4タイプでは動かない。

参考
Google Colab で OpenCALM-7B を試す