Google Colab で AI BunChoで利用されているモデル japanese-novel-gpt-j-6b-v1.5 を試す
「Google Colab」で「AI BunCho」で利用されているモデル「japanese-novel-gpt-j-6b」を試したので、まとめました。
【注意】T4のハイメモリで動作確認しています。
1. japanese-novel-gpt-j-6b
「japanese-novel-gpt-j-6b」は、日本初のAI(大規模言語モデル)を使用した創作支援アプリ「AI BunCho」で利用されているモデルです。「GPT-J-6B」をTPUで2週間日本語トークナイザーを用いて日本語データで事前学習し、その後2週間小説データで転移学習したものになります。
3. Colabでの実行
Colabでの実行手順は、次のとおりです。
(1) Colabのノートブックを開き、メニュー「編集 → ノートブックの設定」で「GPU」の「T4」の「ハイメモリ」を選択。
(2) パッケージのインストール。
# パッケージのインストール
!pip install transformers sentencepiece accelerate(3) トークナイザーとモデルの準備。
from transformers import GPTJForCausalLM, AlbertTokenizer
import torch
# トークナイザーとモデルの準備
tokenizer = AlbertTokenizer.from_pretrained(
"AIBunCho/japanese-novel-gpt-j-6b",
keep_accents=True,
remove_space=False
)
model = GPTJForCausalLM.from_pretrained(
"AIBunCho/japanese-novel-gpt-j-6b",
torch_dtype=torch.float16,
low_cpu_mem_usage=True
)
model.half()
model.eval()
if torch.cuda.is_available():
model = model.to("cuda")(4) 推論の実行。
入力プロンプトの書式は、次のとおりです。
# プロンプトの準備
prompt = "わたくしといふ現象は"
# 推論の実行
input_ids = tokenizer.encode(
prompt,
add_special_tokens=False,
return_tensors="pt"
).cuda()
tokens = model.generate(
input_ids.to(device=model.device),
max_new_tokens=32,
temperature=0.6,
top_p=0.9,
repetition_penalty=1.2,
do_sample=True,
pad_token_id=tokenizer.pad_token_id,
bos_token_id=tokenizer.bos_token_id,
eos_token_id=tokenizer.eos_token_id
)
out = tokenizer.decode(tokens[0], skip_special_tokens=True)
print(out)わたくしといふ現象は、その性質上、非常に強力な魔力を持つ存在に惹かれます。