Google Colab で MPT-7B を試す

「Google Colab」で「MPT-7B」を試したので、まとめました。

【注意】「MPT-7B」を実行するには、「Google Colab Pro/Pro+」のプレミアムが必要です。

1. MPT-7B

「MPT-7B」は、「MosaicML」が開発した、テキストとコードの1Tトークンでゼロから学習したLLMです。オープンソースで、商用利用可能で、LLAMA-7Bと同等の品質を持ちます。約20万ドル、人間介入なし、9.5日間、MosaicMLプラットフォームで学習しています。

Introducing MPT-7B: A New Standard for Open-Source, Commercially Usable LLMs

現在、以下のモデルが提供されています。

・MPT-7B (ベースモデル)
・MPT-7B-StoryWriter-65k+
・MPT-7B-Instruct
・MPT-7B-Chat

2. Colabでの実行

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

(1) メニュー「編集→ノートブックの設定」で、「ハードウェアアクセラレータ」で「GPU」で「プレミアム」を選択。

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

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

(3) トークナイザーとモデルの準備。
今回は、7Bのチャットモデルを指定しています。

import torch
import transformers
from transformers import AutoTokenizer, AutoModelForCausalLM

# トークナイザーとモデルの準備
tokenizer = AutoTokenizer.from_pretrained(
    "mosaicml/mpt-7b-chat"
)
model = AutoModelForCausalLM.from_pretrained(
    "mosaicml/mpt-7b-chat", 
    torch_dtype=torch.float16,
    trust_remote_code=True
).to("cuda:0")

mosaicml/mpt-7b-chat · Hugging Face

(4) 推論の実行。

# プロンプトの準備
prompt = "<human>: Who is Hayao Miyazaki?\n<bot>:"

# 推論の実行
inputs = tokenizer(prompt, return_tensors='pt').to(model.device)
input_length = inputs.input_ids.shape[1]
outputs = model.generate(
    **inputs, 
    max_new_tokens=128, 
    do_sample=True, 
    temperature=0.7, 
    top_p=0.7, 
    top_k=50, 
    return_dict_in_generate=True
)
token = outputs.sequences[0, input_length:]
output_str = tokenizer.decode(token)

# 確認
print("output :", output_str)

output : Hayao Miyazaki is a Japanese film director, animator, and author. He is widely regarded as one of the greatest animators of all time, and is particularly known for his fantasy films.

【翻訳】
output : 宮崎駿は、日本の映画監督、アニメーター、作家です。 彼は史上最高のアニメーターの 1 人として広く認められており、特にファンタジー映画で知られています。

チャットモデルのプロンプトの書式は、RedPajamaと同じものを使ってみました。(学習時に使った書式を後で調べる)

<human>: [Instruction]
<bot>:

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