Google Colab で Youri-7B を試す

「Google Colab」で「Youri-7B」を試したので、まとめました。

1. Youri-7B

「Youri-7B」は、「Rinna」が開発した「Llama 2」ベースの日本語LLMです。

rinna、Llama 2の日本語継続事前学習モデル「Youri 7B」を公開｜rinna株式会社

2. Youri-7B のモデル

「Youri-7B」は、現在6つのモデルが提供されています。

・rinna/youri-7b： 汎用言語モデル
・rinna/youri-7b-instruction : 一問一答の対話モデル
・rinna/youri-7b-chat : 複数ターンの対話モデル
・rinna/youri-7b-gptq： 汎用言語モデル (GPTQ版)
・rinna/youri-7b-instruction-gptq : 一問一答の対話モデル (GPTQ版)
・rinna/youri-7b-chat-gptq : 複数ターンの対話モデル (GPTQ版)

3. Colabでの実行

Colabでの実行手順は、次のとおりです。

(1) Colabのノートブックを開き、メニュー「編集 → ノートブックの設定」で「GPU」を選択。

(2) パッケージのインストール。

# パッケージのインストール
!pip install transformers accelerate bitsandbytes auto_gptq

(2) トークナイザーとモデルの準備。
今回は、複数ターンの対話モデル (GPTQ版)である「rinna/youri-7b-chat-gptq」を使います。

import torch
from transformers import AutoTokenizer
from auto_gptq import AutoGPTQForCausalLM

# トークナイザーとモデルの準備
tokenizer = AutoTokenizer.from_pretrained(
    "rinna/youri-7b-chat-gptq"
)
model = AutoGPTQForCausalLM.from_quantized(
    "rinna/youri-7b-chat-gptq", 
    device_map="auto",
    use_safetensors=True
)

(3) 推論実行の関数の準備。

# 推論の実行
def generate(prompt):
    token_ids = tokenizer.encode(
        prompt, 
        add_special_tokens=False, 
        return_tensors="pt")
    with torch.no_grad():
        output_ids = model.generate(
            input_ids=token_ids.to(model.device),
            max_new_tokens=200,
            do_sample=True,
            temperature=0.5,
            pad_token_id=tokenizer.pad_token_id,
            bos_token_id=tokenizer.bos_token_id,
            eos_token_id=tokenizer.eos_token_id
        )
    return tokenizer.decode(
        output_ids[0][token_ids.size(1) :], 
        skip_special_tokens=True
    )

(4) 推論の実行。

# プロンプトの準備
prompt = """設定: あなたの優秀なAIアシスタントです。
ユーザー: まどか☆マギカで一番かわいいのは？
システム: """

# 推論の実行
print(generate(prompt))

私は、悪魔の少女であるキュゥべえが一番かわいいと思います。

(5) 続きの推論の実行。

# プロンプトの準備
prompt = """設定: あなたの優秀なAIアシスタントです。
ユーザー: まどか☆マギカで一番かわいいのは？
システム: 私は、悪魔の少女であるキュゥべえが一番かわいいと思います。
ユーザー: そのキャラクターの特徴は？
システム: """

# 推論の実行
print(generate(prompt))

彼女は、願いを叶えるために魔法少女たちに取引を持ちかける悪魔です。

OpenAI GPT-4/ChatGPT/LangChain 人工知能プログラミング実践入門