Google Colab で GPT-NEOX-20B による日本語テキスト生成を試す
Google Colabで「GPT-NEOX-20B」による日本語テキスト生成を試したのでまとめました。
1. GPT-NEOX-20B
「GPT-NEOX-20B」は、EleutherAIが提供するオープンソースの200億パラメータの言語モデルです。Pileデータセットで学習されています。
2. Colabでの実行
Google Colabでの実行手順は、次のとおりです。
(1) 新規のColabのノートブックを開き、メニュー「編集 → ノートブックの設定」で「GPU」の「プレミアム」を選択
(2) Googleドライブのマウント。
# Googleドライブのマウント
from google.colab import drive
drive.mount("/content/drive")
(3) 作業フォルダの作成と移動。
マイドライブ直下の「work」に作成しています。
# 作業フォルダの作成と移動
import os
os.makedirs("/content/drive/My Drive/work", exist_ok=True)
%cd "/content/drive/My Drive/work"
(4) モデルのダウンロード。
43GBほどあるため、時間がかかります。
# モデルのダウンロード
!git lfs install
!git clone https://huggingface.co/EleutherAI/gpt-neox-20b
(5) パッケージのインストール。
# パッケージのインストール
!pip install transformers accelerate bitsandbytes
(6) トークナイザーとモデルの準備。
AutoModelForCausalLM.from_pretrained()で、device_map="auto" と load_in_8bit=True を指定します。
from transformers import AutoTokenizer, AutoModelForCausalLM
# トークナイザーの準備
tokenizer = AutoTokenizer.from_pretrained("./gpt-neox-20b")
# モデルの準備
model = AutoModelForCausalLM.from_pretrained(
"./gpt-neox-20b",
device_map="auto",
load_in_8bit=True,
)
(7) 推論の実行。
推論に40秒ほどかかりました。
# プロンプト
prompt = """ずんだもんは東北に住む活発でかわいい女の子です。
あなた「きみの名前は?」
ずんだもん「ボクの名前はずんだもんなのだ。」
あなた「今日の天気は?」
ずんだもん「ボクは今日の天気は晴れだと思うのだ」
あなた「今日は何時に起きたの?」
ずんだもん「ボクは7時に起きたのだ」
あなた「年齢は?」
ずんだもん「ボクの年齢は秘密なのだ」
あなた「昼ごはんは何食べたい?」
ずんだもん「"""
# 推論
input_ids = tokenizer(prompt, return_tensors="pt").input_ids
input_ids = input_ids.to('cuda')
output_tokens = model.generate(
input_ids,
do_sample=True,
max_length=400, # 最長の文章長
min_length=100, # 最短の文章長
temperature=0.9, # 温度
num_return_sequences=1, # 生成する文章の数
)
output = tokenizer.batch_decode(output_tokens)[0]
print(output)
ずんだもんは東北に住む活発でかわいい女の子です。
あなた「きみの名前は?」
ずんだもん「ボクの名前はずんだもんなのだ。」
あなた「今日の天気は?」
ずんだもん「ボクは今日の天気は晴れだと思うのだ」
あなた「今日は何時に起きたの?」
ずんだもん「ボクは7時に起きたのだ」
あなた「年齢は?」
ずんだもん「ボクの年齢は秘密なのだ」
あなた「昼ごはんは何食べたい?」
ずんだもん「昼ごはんはコロッケだと思うのだ。」
RAMの消費は、次のとおりです。
参考
この記事が気に入ったらサポートをしてみませんか?