N番煎じでDeepreneur社のblue-lizardをDatabricksで動かす

久しぶりに日本語モデルをいじってみる。

導入

東大松尾研究室発のAIスタートアップ、Deepreneur社がLlama2をベースとした日本語LLMであるblue-lizardを公開しました。

blue-lizardは70億パラメータと比較的軽量なモデルですが、JGLUE（日本語タスクにおける評価ベンチマーク）を用いた評価では、ChatGPT-3.5を超えるスコアが算出しているようです。
詳細は上記リンク先のベンチマーク結果を参照ください。

なお、Llama2をベースとしているため、ライセンスはLlama2ライセンスに従います。

というわけで、既に試されている方が多数いらっしゃいそうなので、n番煎じシリーズとして試してみます。

検証環境はDatabricks on AWS、DBRは14.3ML LTS、クラスタタイプはg5.xlarge(AWS)を使用しています。

Step1. パッケージインストール

今回は最近お気に入りのSGLangを使って推論します。

パッケージは事前にダウンロードしておいたWheelファイルからインストールしました。
(もちろん直接pipリポジトリからもインストールできます)

# torch, xformers
# pytorchのリポジトリから直接インストールする場合
# %pip install -U https://download.pytorch.org/whl/cu118/torch-2.1.2%2Bcu118-cp310-cp310-linux_x86_64.whl
# %pip install -U https://download.pytorch.org/whl/cu118/xformers-0.0.23.post1%2Bcu118-cp310-cp310-manylinux2014_x86_64.whl
%pip install -U /Volumes/training/llm/tmp/torch-2.1.2+cu118-cp310-cp310-linux_x86_64.whl
%pip install -U /Volumes/training/llm/tmp/xformers-0.0.23.post1+cu118-cp310-cp310-manylinux2014_x86_64.whl

# vLLM
# GithubからvLLMを直接インストールする場合
# %pip install https://github.com/vllm-project/vllm/releases/download/v0.3.0/vllm-0.2.7+cu118-cp310-cp310-manylinux1_x86_64.whl
%pip install /Volumes/training/llm/tmp/vllm-0.3.0+cu118-cp310-cp310-manylinux1_x86_64.whl

# SGLang
%pip install "sglang[srt] >= 0.1.12"

# Triton
# %pip install -U "triton>=2.2.0"
%pip install -U /Volumes/training/llm/tmp/triton-2.2.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

dbutils.library.restartPython()

Step2. モデルダウンロード

以下のリポジトリから、blue-lizardをダウンロード。

その後、UnityCatalog Volumesに保管しておきます。

from typing import Optional

def download_model(model_id:str, revision:Optional[str]=None):
    import os
    from huggingface_hub import snapshot_download

    UC_VOLUME = "/Volumes/training/llm/model_snapshots"

    rev_dir = ("--" + revision) if revision else ""
    local_dir = f"/tmp/{model_id}{rev_dir}"
    uc_dir = f"/models--{model_id.replace('/', '--')}"
    
    snapshot_location = snapshot_download(
        repo_id=model_id,
        revision=revision,
        local_dir=local_dir,
        local_dir_use_symlinks=False,
    )

    dbutils.fs.cp(f"file:{local_dir}", f"{UC_VOLUME}{uc_dir}{rev_dir}", recurse=True)

model_id = "Deepreneur/blue-lizard"
download_model(model_id)

Step3. モデルのロード

torchのmultiprocessingの設定をした上で、モデルをロードします。

import torch
torch.multiprocessing.set_start_method('spawn', force=True)

from sglang import function, system, user, assistant, gen, set_default_backend, Runtime
from sglang.lang.chat_template import get_chat_template

model_path = "/Volumes/training/llm/model_snapshots/models--Deepreneur--blue-lizard"

runtime = Runtime(model_path)

runtime.endpoint.chat_template = get_chat_template("llama-2-chat")

set_default_backend(runtime)

Step4. 推論

準備が整ったので、バッチ推論でいくつか実行してみます。

import time

max_tokens = 256

@function
def simple_question(s, question):
    s += system("あなたは誠実で優秀な日本人のアシスタントです。")
    s += user(question)
    s += assistant(gen("answer", max_tokens=max_tokens))


tic = time.time()

states = simple_question.run_batch(
    [
        {"question": "Databricksとは何ですか？"},
        {"question": "まどか☆マギカでは誰が一番かわいい?"},
        {"question": "ランダムな10個の要素からなるリストを作成してソートするコードをPythonで書いてください。"},
        {"question": "架空の生き物をデザインし、その特徴や能力について説明してください。"},        
        {"question": "日本国内で観光に行きたいと思っています。東京、名古屋、大阪、京都、福岡の特徴を表にまとめてください。列名は「都道府県」「おすすめスポット」「おすすめグルメ」にしてください。"},                
    ],
    temperature=0,
    progress_bar=True,
)

# 結果の出力
for s in states:
    # Message
    for m in s.messages():
        print(m["role"], ":", m["content"])    
    print()
    print("-"*20)
    print()

latency = time.time() - tic

print(f"Latency: {latency:.3f}")

出力

system : あなたは誠実で優秀な日本人のアシスタントです。
user : Databricksとは何ですか？
assistant : Databricksは、Apache Sparkの上に構築されたデータウェアハウスソリューションで、データの探索、作成、保存、そしてインタラクティブな洞察得得を支援することができます。

--------------------

system : あなたは誠実で優秀な日本人のアシスタントです。
user : まどか☆マギカでは誰が一番かわいい?
assistant : 杖のおばあちゃん

--------------------

system : あなたは誠実で優秀な日本人のアシスタントです。
user : ランダムな10個の要素からなるリストを作成してソートするコードをPythonで書いてください。
assistant : 'def sort_list(list): 
  random.shuffle(list) 
  return list
  
list = [5, 2, 8, 7, 10, 4, 6, 3, 1, 9] 

list = sort_list(list) 

print(list) 

# 出力: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]'

--------------------

system : あなたは誠実で優秀な日本人のアシスタントです。
user : 架空の生き物をデザインし、その特徴や能力について説明してください。
assistant : タイフロンは、細身の長い体、2つの翼、そして竜のような頭を持つ神話上の生き物です。その鱗は虹色に輝き、目は明るい黄色です。強力な翼を持っているため、高速で空を飛び、遠くから優れた視力を持って観察することができます。尾は長く、爪は鋭く引っ込められ、捕食者を威嚇することができます。タイフロンは、強力な呪文を使って自分自身を守ることができます。

--------------------

system : あなたは誠実で優秀な日本人のアシスタントです。
user : 日本国内で観光に行きたいと思っています。東京、名古屋、大阪、京都、福岡の特徴を表にまとめてください。列名は「都道府県」「おすすめスポット」「おすすめグルメ」にしてください。
assistant : 東京スカイツリー、名古屋の金のしゃちほこ、大阪の明治神宮、京都の朱雀大路、福岡の熊本城を訪れるのは、日本の観光客にとってすべてメインアトラクションとなる。これらの都市の間を移動するには、日本の鉄道網の幅広い路線網を利用することになる。

--------------------

Latency: 6.876

Step5. 終了処理

最後にラインタイムを終了。

runtime.shutdown()

まとめ

Deepreneur社のblue-lizardをDatabricks上で動かしてみました。

日本語の回答能力は確かに高いと思いました。
用途によってはいい選択肢になるかもしれません。

日本語LLMもかなり増えてきました。選択肢ができるのは非常によいと思うので、この調子でいろんなモデルがリリースされていくとありがたいなあ。