rinna社が日本語特化の対話GPT言語モデルを公開！ VRAM6GBあれば動くかも。

prompt = [
    {
        "speaker": "ユーザー",
        "text": "最近、眠りが浅くい気がして改善できないかなぁと思っています。"
    },
    {
        "speaker": "システム",
        "text": "もう少し具体的に説明していただけますか？"
    },
    {
        "speaker": "ユーザー",
        "text": "すっきり起きられない気がします。日中も眠たくなります。"
    },
    {
        "speaker": "システム",
        "text": "分かりました、十分な睡眠を実感できていないということですね。睡眠時間はどれくらいですか？"
    },
    {
        "speaker": "ユーザー",
        "text": "６時間ぐらいなのですが...。どうすれば改善できるでしょうか。"
    }
]
prompt = [
    f"{uttr['speaker']}: {uttr['text']}"
    for uttr in prompt
]
prompt = "<NL>".join(prompt)
prompt = (
    prompt
    + "<NL>"
    + "システム: "
)

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("rinna/japanese-gpt-neox-3.6b-instruction-ppo", use_fast=False)
model = AutoModelForCausalLM.from_pretrained("rinna/japanese-gpt-neox-3.6b-instruction-ppo")

if torch.cuda.is_available():
    model = model.to("cuda")

token_ids = tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt")

with torch.no_grad():
    output_ids = model.generate(
        token_ids.to(model.device),
        do_sample=True,
        max_new_tokens=128,
        temperature=0.7,
        repetition_penalty=1.1,
        pad_token_id=tokenizer.pad_token_id,
        bos_token_id=tokenizer.bos_token_id,
        eos_token_id=tokenizer.eos_token_id
    )

output = tokenizer.decode(output_ids.tolist()[0][token_ids.size(1):])
output = output.replace("<NL>", "\n")
print(output)

model = AutoModelForCausalLM.from_pretrained(
    "rinna/japanese-gpt-neox-3.6b-instruction-ppo",
    device_map='auto',
    torch_dtype=torch.float16
    )

model = AutoModelForCausalLM.from_pretrained(
    "rinna/japanese-gpt-neox-3.6b-instruction-ppo",
    device_map={"": 0},
    # torch_dtype=torch.float16,
    load_in_4bit=True
    )

# if torch.cuda.is_available():
#     model = model.to("cuda")