MacBook Proでllm-jp-3-13b-instruct-4bitをMLXと共に試してみる

2024年9月26日 10:57

「オープンかつ日本語に強いGPT-3級大規模言語モデル「LLM-jp-3 172B」の事前学習に使用しているコーパスで学習した 1.8B、3.7B、13B の各モデルを公開」したらしく、そのMLX変換モデルが公開されていましたので、試してみます。

今回は、13B-instructモデルの4bit量子化モデルを試します。

環境は、MacBook Pro (Apple M3 Pro) です。

1. 準備

仮想環境

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

つづいて、パッケージのインストール。

pip install -U mlx-lm

2. 流し込むコード

こちらを/path/to/query4mlx.pyとして保存します。

import sys
import argparse
from mlx_lm import load, generate, stream_generate
from typing import List, Dict
import time

# argv
parser = argparse.ArgumentParser()
parser.add_argument("--model-path", type=str, default=None)
parser.add_argument("--tokenizer-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=4096)

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_tokens = args.max_tokens

model, tokenizer = load(path_or_hf_repo=model_id)

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

def q(
    user_query: str,
    history: List[Dict[str, str]]=None
) -> List[Dict[str, str]]:
    # generation params
    generation_params = {
        "temp": 0.8,
        "top_p": 0.95,
        "max_tokens": max_tokens,
        "repetition_penalty": 1.1,
        "repetition_context_size": 20,
    }
    #
    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=""
    for t in stream_generate(
        model=model,
        tokenizer=tokenizer,
        prompt=prompt,
        **generation_params
    ):
        print(t, end="", flush=True)
        output += t
    print()
    output_ids = tokenizer.encode(
        output,
        add_special_tokens=True,
        return_tensors="pt"
    )
    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])
    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. 試してみる

起動します。

python -i ~/scripts/query4mlx.py --model-path mlx-community/llm-jp-3-13b-instruct-4bit

聞いてみましょう。

>>> history = q("ドラえもんとはなにか")
--- prompt
<s>以下は、タスクを説明する指示です。要求を適切に満たす応答を書きなさい。

### 指示:
ドラえもんとはなにか

### 応答:

--- output

「ドラえもん」とは、藤子・F・不二雄による日本の漫画、およびそれを原作としたアニメ作品です。
22世紀の未来からやってきたネコ型ロボット「ドラえもん」と、勉強もスポーツも苦手な小学生・野比のび太の日常生活を描いた作品で、夢いっぱいのSFの世界と友情、努力、家族愛といったテーマが描かれています。
以下に、ドラえもんについて詳しく説明します。

1. 概要
「ドラえもん」は、小学館の雑誌『よいこ』、『幼稚園』、および『小学一年生』1970年1月号から1974年5月号の間に連載された作品です。
てんとう虫コミックスでは、全45巻が刊行されており、この作品は日本を代表する国民的漫画の一つとなっています。
連載は、1969年1月号から1974年5月号までと長期に渡り、一度は終了しましたが、1979年4月2日よりテレビアニメ化されて以降も現在まで30年以上にわたって放送されており、全国に知られるようになりました。
また、日本国内外で多くの賞を受賞しており、原作は全世界でも30カ国語以上に翻訳され、海外でも広く知られています。

2. ストーリー
22世紀の未来からやってきたネコ型ロボット「ドラえもん」と、勉強もスポーツも苦手な小学生・野比のび太の日常生活を描いた作品です。
のび太が困ったときに、ドラえもんから与えられるひみつ道具でさまざまな問題を解決しようとします。しかし、それで大抵は失敗し、結局、のび太とドラえもんが一緒になって困った状況に立ち向かうというストーリーです。
のび太は「クラスのガキ大将」として描かれ、勉強もスポーツも苦手ですが、優しく、正義感が強くて友達思いの男の子として描かれています。
作品には、夢いっぱいのSFの世界と友情、努力、家族愛といったテーマが描かれており、のび太やジャイアン、スネ夫、しずかといった子どもたちが成長していく様子も描かれています。
なお、ドラえもんは、どら焼きが好物で、トレードマークの四次元ポケットには、便利なひみつ道具が多数入っています。

3. キャラクター
ドラえもん:22世紀の未来からやってきたネコ型ロボット。のび太の家に住み、のび太と一緒に様々な冒険を繰り広げます。
野比のび太:勉強もスポーツも苦手で、ぐうたらな性格。「ドラえもん」の主人公です。
源静香:のび太と同じ小学校に通う5年生で、のび太が憧れるクラスメイト。水泳とピアノが得意な女の子です。
剛田武:通称「ジャイアン」で、のび太のクラスメイト。ガキ大将で力持ちですが、仲間思いな一面もあります。
骨川スネ夫:のび太のクラスメイトで、裕福な家庭で育ったお坊ちゃま。自慢話が多いです。
のび太のママ:のんびりした性格で、のび太にいつも優しく接しています。

mlx-community/llm-jp-3-13b-instruct-4bit

prompt tokens = 37
output tokens = 556 (87.817432 [tps])
   total time = 6.331317 [s]

すばらしい。

メモリの使用量は8.63GBとのこと。

MacBook Proでllm-jp-3-13b-instruct-4bitをMLXと共に試してみる

1. 準備

2. 流し込むコード

3. 試してみる

関連