MacでAirLLMで動かす

AI

初めに

MacでLLMで動くライブラリ「AirLLM」があると聞いたので、実際に試してみます

github.com

参考サイト

note.com

環境

準備

環境を構築します

python3 -m venv air_llm
source bin/activate

ライブラリをインストールします

pip install mlx airllm

コードの作成

推論するために、モデルのロードや設定などを書いた main.pyを作成します

from airllm import AutoModel
import mlx.core as mx
import time

llm  = "elyza/ELYZA-japanese-Llama-2-13b-instruct"

# モデルの準備
model = AutoModel.from_pretrained(llm)

# Llama 2 based
B_INST, E_INST = "[INST]", "[/INST]"
B_SYS, E_SYS = "<<SYS>>\n", "\n<</SYS>>\n\n"
DEFAULT_SYSTEM_PROMPT = "あなたは誠実で優秀な日本人のアシスタントです。"

def build_prompt(user_query, chat_history=None):
    prompt = "{chat_history}{b_inst} {system}{prompt} {e_inst} ".format(
        chat_history=chat_history,
        b_inst=B_INST,
        system=f"{B_SYS}{DEFAULT_SYSTEM_PROMPT}{E_SYS}",
        prompt=user_query,
        e_inst=E_INST,
    )
    return prompt

def q(user_query, chat_history=None):
    start = time.process_time()
    # 推論の実行
    prompt = build_prompt(user_query, chat_history)
    input_ids = model.tokenizer(
        model.tokenizer.bos_token + prompt,
        return_tensors="np",
        return_attention_mask=False,
        truncation=True,
        max_length=128,
        padding=False
    )
    output_ids = model.generate(
        mx.array(input_ids['input_ids']),
        max_new_tokens=3,
        use_cache=True,
        return_dict_in_generate=True
    )
    print(output_ids)
    chat_history = prompt + output_ids
    end = time.process_time()
    print(end - start)
    return chat_history

chat_history = q("What is Doraemon?")

上記を 実行します

python main.py

初めは、以下のようにモデルのロードが走ります

(air_llm) user air_llm % python main.py
pytorch_model.bin.index.json: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 29.9k/29.9k [00:00<00:00, 11.3MB/s]
README.md: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4.64k/4.64k [00:00<00:00, 2.07MB/s]
config.json: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 695/695 [00:00<00:00, 496kB/s]
.gitattributes: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1.52k/1.52k [00:00<00:00, 5.54MB/s]
key_visual.png: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 277k/277k [00:00<00:00, 10.5MB/s]
generation_config.json: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 132/132 [00:00<00:00, 536kB/s]
special_tokens_map.json: 100%|███████

推論

推論のプロンプトは 、以下を使用します.

What is Doraemon?

結果

以下の部分を変更してmax new tokensを3及び100の両方で推論を行なってみました

max_new_tokens=3,

tokens = 3の場合

Doraemon
24.876844

tokens = 100の場合

Doraemon (ドラえもん) is a popular Japanese manga and anime series created by Fujiko Fujio. It follows the adventures of a robotic cat named Doraemon, who travels back in time from the 22nd century to help a young boy named Noby (のび太) and his friends with various problems and adventures. The series is known for its blend of humor, adventure, and science fiction elements, and
905.053268

まとめ

推論速度はかなり遅いので動くのは動きそうですが、実際に使うのは難しそうです