見出し画像

Google Colab で AI BunChoで利用されているモデル japanese-novel-gpt-j-6b-v1.5 を試す

npaka

「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)
わたくしといふ現象は、その性質上、非常に強力な魔力を持つ存在に惹かれます。

関連



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