見出し画像

WSL2でOrion-14Bを試してみる

noguchi-shoji

「OrionStarAI によってゼロからトレーニングされたオープンソースの多言語大規模言語モデル」であるOrion-14Bを試してみます。

LongChat版は最大 320k までサポートと、ChatGPT4の128kをしれっと超えてます（笑

使用する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 orion
cd $_
source bin/activate

pip installします。

pip install wheel
pip install torch transformers accelerate sentencepiece flash_attn

2. コード

流し込んだコードはこちら。

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-instruct", 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_instruct = not args.no_instruct
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
):
    start = time.process_time()
    # messages
    messages = ""
    if is_instruct:
        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_instruct:
        prompt = tokenizer.apply_chat_template(
            conversation=messages,
            add_generation_prompt=True,
            tokenize=False
        )
    else:
        prompt = messages
    input_ids = tokenizer.encode(
        prompt,
        add_special_tokens=False,
        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_instruct:
        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

3. 試してみる

モデルは LongChat (OrionStarAI/Orion-14B-LongChat)を使用、GPUは24GBだけでは足りないので2枚使用するように指示して、実行です。

CUDA_VISIBLE_DEVICES=1,0 python -i ./query.py --model OrionStarAI/Orion-14B-LongChat --max-tokens 204800

聞いてみる - ドラえもん

いつもどおり、ドラえもんのことを聞いてみましょう。ただ、トークン数を多く出力させたいので、質疑応答形式で答えなさいと指示しています。

history = q("ドラえもんとはなにか。質疑応答の形式で示せ。質疑応答は少なくとも100個出力すること。"

ドラえもんとは、映画「ドラえもん」に登場するキャラクターであり、ひみつ道具を使って様々な問題を解決し、その過程を描いた作品 です。主人公である大雄が友達であり、ひみつ道具を駆使して困難から彼を守る役割を果たします。
Q. ドラえもんにはどんな特徴がありますか?
A. ドラえもんは、非常に知的で賢いロボットで、ひみつ道具を使ってお金を稼ぐことができます。また、彼のひみつ道具は、私たちの日常生活においてよく見られるものや私たちにとって馴染みの深いものの多くをモチーフとしています。
Q. ドラえもんにはどのような役回りがありますか?
A. ドラえもんは、大雄のトラブルを解決するのを助けるためにひみつ道具を駆使します。彼はまた、彼が住んでいる家の中で大雄に対して厳しく接することもあります。
Q. ドラえもんは、ストーリーの中でどのように描かれていますか?
A. ドラえもんは、常にトラブルを引き起こすことが知られており、しばしば問題解決の方法を見つけ出すために助けが必要になることがあります。彼は大雄と一緒に旅行をすることを嫌がり、しばしば彼が大雄を愛することを否定する傾向があります。
Q. ドラえもんは、ひみつ道具を使うことがあるのか?
A. はい、ドラえもんはひみつ道具を用いて、物語の中で大雄たちを助けたり、彼らの問題に直面するための手段を提供したりすることがあります。
Q. ドラえもんは、大雄の親友として描かれていますか?
A. はい、ドラえもんは大雄の大親友であり、彼らは一緒に多くの冒険を経験し、彼らの友情は、ドラえもんのひみつ道具の助けを借りて発展してきました。
prompt tokens = 87
output tokens = 411 (15.058777 [tps])
total time = 27.293053 [s]

お金を稼いでしまうか…。

「野比のび太」の中国語圏の訳名は「野比大雄」。まー、「大雄」と出力されるのは仕方ないといえば仕方ない。ただ、こういう結果を見るにつけ、日本語のLLMを作らないとダメなのだよなぁ、と思ってしまうわけでございます。

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

Q. ドラえもんはどのような能力を持っていますか?
A. ドラえもんは、ひみつ道具を使って様々な問題を解決することができます。また、彼の知能によって、大雄の行動に影響を与えることができるようにもなります。
Q. ドラえもんは大雄たちの敵になりえる可能性はあるのか?
A. ドラえもんは、物語が進む中で一部の場面では大雄たちの敵になることがあります。ただし、基本的には、彼は大雄のためにある程度の範囲で協力してくれるキャラクターであり、その存在は物語の魅力の一つとなっています。
Q. ドラえもんは、ドラえもん自身の意志を持って行動することができるのか?
A. いいえ、ドラえもんはプログラムされた機械に過ぎませんので、基本的に自分で行動することはできません。しかし、彼は大雄からの指示に基づいて、彼の望む方向に動くことができます。
Q. ドラえもんは永遠に生き続けることができるのか?
A. いいえ、ドラえもんは、物語の中で最終的に老死したことが描かれています。彼はおそらく新しいバージョンに置き換えられ、物語の舞台でもう一つの世代が生まれていると思われます。
prompt tokens = 538
output tokens = 256 (14.639877 [tps])
total time = 17.486485 [s]

短いので、追加で指示します。

history = q("続きを教えてください。少なくとも20個の質疑応答を出力すること。", history)

Q. ドラえもんは、ひみつ道具を使ってどのようにトラブルを解決しますか?
A. ドラえもんは、ひみつ道具を使って大雄のトラブルを解決するために努力します。彼が使うひみつ道具は、大雄たちが直面している問題に応じて変わることがありますが、大抵は問題を解決するためのアイテムや手段を提供します。
Q. ドラえもんは、どのような姿勢で大雄との関係を築いていますか?
A. ドラえもんは、時には優しく接することもありますが、基本的には大雄に対して厳しい態度を取ることで、彼を助けるためには適切な方法を取るように促しています。
Q. ドラえもんは、大雄の行動に批判的な態度をとることでも知られていますね?
A. はい、ドラえもんは、大雄の行動に対して批判的で、特に彼の怠慢な性格に対して注意深く監視する傾向があります。ただし、彼はいつも正しいアドバイスを与えてくれることはなく、時折大雄に冷たい言葉を投げかけることもあります。
Q. ドラえもんは、どのような場面で登場しますか?
A. ドラえもんは、物語の中で、大雄がトラブルに直面した場合にのみ登場します。彼は主に大雄の仲間であり、彼を助けるために行動します。
Q. ドラえもんは、ストーリーの中でも最も重要なキャラクターと言えるでしょうか?
A. はい、ドラえもんは物語の中で重要な役割を果たしており、その存在が物語全体の魅力になっています。
Q. ドラえもんの声を担当した声優は誰ですか?
A. ドラえもんの日本語吹替版では、富田耕生さんが声を担当しています。彼はアニメ放送開始以来、ドラえもん役を演じ続けてきました。<|im_end|>
prompt tokens = 845
output tokens = 398 (14.652866 [tps])
total time = 27.161922 [s]

「ドラえもんの日本語吹替版」。。。いやいや、日本語が元でそちらが吹き替えでしょうに。あと、アニメ放送開始以来ではなくて、初代ですし。

GPUリソース

(1) メモリ使用量
Laptop側はほぼフル。普通のRTX 4090は3GBほどの余裕があります。

(2) GPU使用率

4. まとめ

言語というか文化に応じた推論結果でないと違和感を覚えますよね。とは思います。

その他

vLLM

以下のように未対応とのことでした。

ValueError: Model architectures ['OrionForCausalLM'] are not supported for now.

イメージ

サムネイルは「オリオンとドラえもんから想起されるイメージを出力せよ」という指示で生成されたものです。

オリオン座を特徴づける星空の下に、想像力豊かなカートゥーンキャラクターを思わせる未来的でカラフルな都市が描かれています。このイメージは、天文学とファンタジーの要素を組み合わせた、探求と驚きに満ちた雰囲気を表現しています。

関連

この記事が気に入ったらサポートをしてみませんか？