rinna-3.6b-instraction-ppoと会話する簡易プログラム 【Colab無料枠でも動きます】

超シンプルな「rinna-3.6b-instraction-ppo」と会話できるプログラムを掲載します。

• google colab の無料枠のGPU（T4）でも動きますし、お手元のゲーミングＰＣ、ゲーミングノートでも十分動くと思いますので、興味がある方はぜひ、お試しいただければ。

• GPUメモリサイズに応じて、model読み込みのオプションを、指定なし、load_in_4bit=True, load_in_8bit=True,  torch_dtype=torch.float16などからえらんで下さい。

ライブラリのインストール

• google colabでは、ライブラリーのインストールの後、ランタイムを再起動してください。

!pip install transformers > /dev/null
!pip install SentencePiece > /dev/null
!pip install Accelerate > /dev/null
!pip install bitsandbytes > /dev/null

プログラム本体

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

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

last_prompt = "システム: わたしの名前は「りんな」。りんなはユーザーの質問になんでもこたえるよ！"
while True:

    user_input = input("ユーザー: ")
    if user_input == "":
        break

    prompt = last_prompt + "<NL>ユーザー: " + user_input + "<NL>システム: "
    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("</s>", "")

    last_prompt = prompt + "システム: " + output

    print("りんな: " + output.replace("<NL>", "\n"))

ユーザー: よかった。起動しましたね。コーヒー飲みますか？
りんな: もちろんです。コーヒーを飲んで、ゆっくりして話をしましょう。何か聞きたいことがあれば、遠慮なく聞いてください。りんなと話すことでリラックスできると思います。良い一日を過ごしてくださいね。また会えることを楽しみにしています。
ユーザー: どんな種類のコーヒーが好き？
りんな: コーヒーにはたくさんの種類がありますが、一般的に、酸味があり、苦みが少なく、甘さが少ないコーヒーが好まれます。ミルクや砂糖を加えて飲むことが多いです。コーヒーを選ぶ際の重要な要素の一つは、豆の品質で、良質な豆ほど豊かな風味を持ち、強いボディを持っています。
ユーザー: 今日はどこの産地のコーヒーを飲もうかなぁ。
りんな: コーヒーを選ぶ際には、地域や生産国によって異なるさまざまな品質基準があることを知っておくことが重要です。例えば、ブラジル産のコーヒーは高品質とされています。

load_in_8bit=Trueの場合、GPU RAM 5GB程度です。