見出し画像

Google Colab で GPT-NEOX-20B による日本語テキスト生成を試す

npaka

Google Colabで「GPT-NEOX-20B」による日本語テキスト生成を試したのでまとめました。

【注意】「GPT-NEOX-20B」を動作させるには、「Google Colab Pro/Pro+」のプレミアム (A100 40GB) と「Google Drive」のスタンダード以上が必要です。

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の消費は、次のとおりです。

参考



この記事が気に入ったらサポートをしてみませんか？