CyberAgentのモデルを試したところですが、Rinna社も新しいモデルを出しました。
rinna、日本語に特化した36億パラメータのGPT言語モデルを公開｜rinna株式会社のプレスリリース
ここですね。
https://huggingface.co/rinna/japanese-gpt-neox-3.6b
ということで試します。
あれーいろいろ試したときはちゃんと東京、いや大阪って答えて面白かったのだけど、ちゃんと答えてくれない・・・
GPUメモリは14.8GB使っています。
torch_dtype=torch.float16を付けても8GBちょうどくらいのGPUメモリを使うので、8GB VRAMだとつらそうです。12GB VRAMであれば大丈夫そう。
そして、今回はチャット向けにトレーニングされたモデルも用意されています。
https://huggingface.co/rinna/japanese-gpt-neox-3.6b-instruction-sft
ユーザーをユーザー、AIからの応答をシステムとして改行を<NL>にして与えればいいようなので、次のような変換を行います。
prompt = f"ユーザー:{prompt}<NL>システム:"
ちゃんと東京って答えてくれます。
あと、Rinnaモデルはなんか面白い返答を返しがち。
オススメ観光地は答えてくれませんでした。
けど、これまでのチャットモデルは3Bくらいだと続きのユーザー入力システム応答を３往復分くらい生成したり、チャットボットとしての動きが不安定だったけど、このモデルはそこはちゃんとしてそう。
オススメ観光地は答えてくれませんが、ユーザーとシステムのやりとりの枠組みが破綻するものは見かけませんでした。
チャット用スクリプトはこんな感じ
プロンプト生成以外でCyberAgentモデルのときと大きく違うのはTokenizerにuse_fast = Falseがついているくらい。
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from colorama import Fore, Back, Style, init
init(autoreset=True)
model_name = "rinna/japanese-gpt-neox-3.6b-instruction-sft"
print ("model:" + model_name)
# model = AutoModelForCausalLM.from_pretrained(model_name).to("cpu")
# model = AutoModelForCausalLM.from_pretrained(model_name, device_map='auto', torch_dtype=torch.float16)
model = AutoModelForCausalLM.from_pretrained(model_name, device_map='auto')
tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast = False)
# prompt = "AIによって私達の暮らしは?"
# prompt = "日本の首都はどこですか？"
# prompt = "吾輩は猫といえば？"
prompt = "日本のおすすめの観光地を教えてください。"
prompt = f"ユーザー:{prompt}<NL>システム:"
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
with torch.no_grad():
tokens = model.generate(
**inputs,
max_new_tokens=100,
do_sample=True,
temperature=0.8,
pad_token_id=tokenizer.pad_token_id,
)
prompt = prompt.replace("<NL>", "\n")
output = tokenizer.decode(tokens[0], skip_special_tokens=True).replace("<NL>", "\n")
print(f"{Fore.YELLOW}{prompt}{Fore.WHITE}{output[len(prompt):]}")
