Google Colab で Japanese Stable LM 3B-4E1T を試す

「Google Colab」で「Japanese Stable LM 3B-4E1T」を試したので、まとめました。

1. Japanese Stable LM 3B-4E1T

「Japanese Stable LM 3B-4E1T」は、「Stability AI」が開発した「Stable LM 3B-4E1T」ベースの日本語LLMです。

日本語大規模言語モデル「Japanese Stable LM 3B-4E1T」「Japanese Stable LM Gamma 7B」を公開しました — Stability AI Japan

2. Japanese Stable LM 3B-4E1T のモデル

「Japanese Stable LM 3B-4E1T」は、現在2つのモデルが提供されています。

・Japanese Stable LM 3B-4E1T Base : ベースモデル
・Japanese Stable LM 3B-4E1T Instruct : 指示モデル

3. Colabでの実行

Colabでの実行手順は、次のとおりです。

(1) Colabのノートブックを開き、メニュー「編集 → ノートブックの設定」で「GPU」を選択。

(2) パッケージのインストール。

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

(2) トークナイザーとモデルの準備。

from transformers import AutoTokenizer, AutoModelForCausalLM

# トークナイザーとモデルの準備
tokenizer = AutoTokenizer.from_pretrained(
    "stabilityai/japanese-stablelm-3b-4e1t-instruct"
)
model = AutoModelForCausalLM.from_pretrained(
    "stabilityai/japanese-stablelm-3b-4e1t-instruct",
    torch_dtype="auto",
    device_map="auto",
    trust_remote_code=True,
)

(3) 推論の実行。

# プロンプトの準備
prompt = """以下は、タスクを説明する指示と、文脈のある入力の組み合わせです。要求を適切に満たす応答を書きなさい。

### 指示: 
まどか☆マギカでは誰が一番かわいい?

### 応答:
"""

# 推論の実行
input_ids = tokenizer.encode(
    prompt, 
    add_special_tokens=False, 
    return_tensors="pt"
)
tokens = model.generate(
    input_ids.to(device=model.device),
    max_new_tokens=100,
    temperature=1,
    top_p=0.3,
    do_sample=True,
)
output = tokenizer.decode(tokens[0][input_ids.shape[1]:], skip_special_tokens=True).strip()
print(output)

まどか☆マギカでは、ほむらが一番かわいいと思います。

OpenAI GPT-4/ChatGPT/LangChain 人工知能プログラミング実践入門