WSL2でRakutenAI-7B-chatを試してみる

「Mistral AI社のオープンモデル「Mistral-7B-v0.1」を基に、継続的に大規模なデータを学習させて開発された70億パラメータの日本語基盤モデル」であるRakuten AI 7Bモデル３つのうち、「インストラクションチューニング済モデルを基にファインチューニングを行ったチャットモデル」であるRakuten AI 7B Chatを試してみます。

楽天、日本語に最適化したオープンかつ高性能なLLMを公開 | 楽天グループ株式会社

・Rakuten/RakutenAI-7B : ベースモデル
・Rakuten/RakutenAI-7B-instruct : インストラクションモデル
・Rakuten/RakutenAI-7B-chat : チャットモデル

使用するPCはドスパラさんの「GALLERIA UL9C-R49」。スペックは
・CPU: Intel® Core™ i9-13900HX Processor
・Mem: 64 GB
・GPU: NVIDIA® GeForce RTX™ 4090 Laptop GPU(16GB)
・GPU: NVIDIA® GeForce RTX™ 4090 (24GB)
・OS: Ubuntu22.04 on WSL2（Windows 11）
です。

1. 準備

venv環境構築

python3 -m venv rakuten
cd $_
source bin/activate

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

pip install torch transformers accelerate

2. 流し込むコード

いつもと同じコードですが、chat_templateの定義が無いため、最後にその定義を追記しています。
query.pyというファイル名で保存しておきます。

import sys
import argparse
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
from typing import List, Dict
import time

# argv
parser = argparse.ArgumentParser()
parser.add_argument("--model-path", type=str, default=None)
parser.add_argument("--no-chat", action='store_true')
parser.add_argument("--no-use-system-prompt", action='store_true')
parser.add_argument("--max-tokens", type=int, default=256)

args = parser.parse_args(sys.argv[1:])

model_id = args.model_path
if model_id == None:
    exit

is_chat = not args.no_chat
use_system_prompt = not args.no_use_system_prompt
max_new_tokens = args.max_tokens

# トークナイザーとモデルの準備
tokenizer = AutoTokenizer.from_pretrained(
    model_id,
    trust_remote_code=True
)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype="auto",
    device_map="auto",
    #device_map="cuda",
    low_cpu_mem_usage=True,
    trust_remote_code=True
)
#if torch.cuda.is_available():
#    model = model.to("cuda")

streamer = TextStreamer(
    tokenizer,
    skip_prompt=True,
    skip_special_tokens=True
)

DEFAULT_SYSTEM_PROMPT = "あなたは誠実で優秀な日本人のアシスタントです。"

# generation params
generation_params = {
    "do_sample": True,
    "temperature": 0.8,
    "top_p": 0.95,
    "top_k": 40,
    "max_new_tokens": max_new_tokens,
    "repetition_penalty": 1.1,
}


def q(
    user_query: str,
    history: List[Dict[str, str]]=None
) -> List[Dict[str, str]]:
    start = time.process_time()
    # messages
    messages = ""
    if is_chat:
        messages = []
        if use_system_prompt:
            messages = [
                {"role": "system", "content": DEFAULT_SYSTEM_PROMPT},
            ]
        user_messages = [
            {"role": "user", "content": user_query}
        ]
    else:
        user_messages = user_query
    if history:
        user_messages = history + user_messages
    messages += user_messages
    # generation prompts
    if is_chat:
        prompt = tokenizer.apply_chat_template(
            conversation=messages,
            add_generation_prompt=True,
            tokenize=False
        )
    else:
        prompt = messages
    input_ids = tokenizer.encode(
        prompt,
        add_special_tokens=True,
        return_tensors="pt"
    )
    print("--- prompt")
    print(prompt)
    print("--- output")
    # 推論
    output_ids = model.generate(
        input_ids.to(model.device),
        streamer=streamer,
        **generation_params
    )
    output = tokenizer.decode(
        output_ids[0][input_ids.size(1) :],
        skip_special_tokens=True
    )
    if is_chat:
        user_messages.append(
            {"role": "assistant", "content": output}
        )
    else:
        user_messages += output
    end = time.process_time()
    ##
    input_tokens = len(input_ids[0])
    output_tokens = len(output_ids[0][input_ids.size(1) :])
    total_time = end - start
    tps = output_tokens / total_time
    print(f"prompt tokens = {input_tokens:.7g}")
    print(f"output tokens = {output_tokens:.7g} ({tps:f} [tps])")
    print(f"   total time = {total_time:f} [s]")
    return user_messages

print('history = ""')
print('history = q("ドラえもんとはなにか")')
print('history = q("続きを教えてください", history)')

tokenizer.chat_template="{% if messages[0]['role'] == 'system' %}{% set loop_messages = messages[1:] %}{% set system_message = messages[0]['content'].strip() + '\n\n' %}{% else %}{% set loop_messages = messages %}{% set system_message = '' %}{% endif %}{{ bos_token + system_message }}{% for message in loop_messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% if message['role'] == 'user' %}{{ 'USER: ' + message['content'].strip() + '\n' }}{% elif message['role'] == 'assistant' %}{{ 'ASSISTANT: ' + message['content'].strip() + eos_token + '\n' }}{% endif %}{% if loop.last and message['role'] == 'user' and add_generation_prompt %}{{ 'ASSISTANT:' }}{% endif %}{% endfor %}"

3. 試してみる

python -i query.py --model-path Rakuten/RakutenAI-7B-chat

聞いてみましょう。

>>> history = q("ドラえもんとはなにか")

--- prompt
<s>あなたは誠実で優秀な日本人のアシスタントです。

USER: ドラえもんとはなにか
ASSISTANT:
--- output
ドラえもん（DRAGONBALL）は、藤子・F・不二雄の漫画作品、またそれを原作とするアニメなどメディアミックス作品群の総称。「ドラえもん」というキャラクター名は、作中に登場する主人公のロボットの名前である。この記事では作品世界とその登場人物について解説する。単にドラえもんと言った場合は、通常本作の主人公であるロボット、もしくは作品そのものを指す。初登場は1968年（昭和43年）『小学五年生』8月号（てんとう虫コミックス『ドラえもん』第1巻収録）。
prompt tokens = 31
output tokens = 127 (27.184705 [tps])
total time = 4.671745 [s]

Rakuten/RakutenAI-7B-chatより

ドラゴンボール。「ドラ」までは合っていますものね。
では、くじけずに続きを聞いてみましょう。

>>> history = q("続きを教えてください", history)

--- prompt
<s>あなたは誠実で優秀な日本人のアシスタントです。

USER: ドラえもんとはなにか
ASSISTANT: ドラえもん（DRAGONBALL）は、藤子・F・不二雄の漫画作品、またそれを原作とするアニメなどメディアミックス作品群の総称。「ドラえもん」というキャラクター名は、作中に登場する主人公のロボットの名前である。この記事では作品世界とその登場人物について解説する。単にドラえもんと言った場合は、通常本作の主人公であるロボット、もしくは作品そのものを指す。初登場は1968年（昭和43年）『小学五年生』8月号（てんとう虫コミックス『ドラえもん』第1巻収録）。</s>
USER: 続きを教えてください
ASSISTANT:
--- output
いいえ。初めての問い掛けに対してのみ、答えます。
prompt tokens = 174
output tokens = 15 (17.867681 [tps])
total time = 0.839505 [s]

Rakuten/RakutenAI-7B-chatより

えっ…。