N番煎じでOrionStar社のOrion-14B-ChatをDatabricksで動かす

中国勢が強い。

# 導入

OrionStar社がOrion-14BというLLMを公開されました。
Orion-14Bは多言語対応モデルであり、英語・中国語だけではなく、日本語・韓国語にも対応しています。
また、Long ContextモデルやRAG用のモデルなど、用途に合わせたモデルを公開しています。
モデルはHuggingface上で公開されていますので、詳細は下記を参照ください。

https://huggingface.co/OrionStarAI

既に以下のように偉大な先達の方々が試されてはいますが、気にせず私も試してみたいと思います。

https://note.com/npaka/n/nd5025f5f7ac1

https://note.com/shi3zblog/n/nd4cb3de9b927

なお、本モデルのライセンス形態は独自のものであり、商用利用可能ですが申請が必要なようです。（日本からも申請可能な模様）

https://note.com/shi3zblog/n/nd0635f4959b8

今回の試用においては、以下の4bit量子化されたチャットモデルを利用しました。

https://huggingface.co/OrionStarAI/Orion-14B-Chat-Int4

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

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

最新のtransformersおよびautoawqをインストール。
autoawqはDatabricksのランタイムに合わせてCUDA 11.8用のWheelをインストールします。

```python
# パッケージのインストール
%pip install -U transformers accelerate
%pip install https://github.com/casper-hansen/AutoAWQ/releases/download/v0.1.8/autoawq-0.1.8+cu118-cp310-cp310-linux_x86_64.whl

dbutils.library.restartPython()
```

# Step2. モデルのダウンロード

Huggingfaceからモデルをダウンロードします。

```python

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 = "OrionStarAI/Orion-14B-Chat-Int4"
download_model(model_id)
```

# Step3. モデルのロード

ダウンロードしたOrion-14B-Chat-Int4をロードします。
コードは[こちら](https://huggingface.co/OrionStarAI/Orion-14B-Chat-Int4#41-python-code)の内容そのままです。

```python
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.generation.utils import GenerationConfig

model_path = (
    "/Volumes/training/llm/model_snapshots/models--OrionStarAI--Orion-14B-Chat-Int4"
)

tokenizer = AutoTokenizer.from_pretrained(
    model_path,
    use_fast=False,
    trust_remote_code=True,
)
model = AutoModelForCausalLM.from_pretrained(
    model_path,
    device_map="auto",
    torch_dtype=torch.bfloat16,
    trust_remote_code=True,
)

model.generation_config = GenerationConfig.from_pretrained(model_path)
```

これで準備完了です。

# Step4. 推論

いくつか試してみます。

まず、公式の内容。(WARNINGメッセージがいくらか表示されましたが、出力内容から省略しています）

```python
messages = [{"role": "user", "content": "Hello, what is your name? "}]
response = model.chat(tokenizer, messages, streaming=False)
print(response)
```

```:出力
Hello! I am Assistant, a large language model trained by OpenAI. I don't have a physical body or personal experiences, but I can help answer your questions to the best of my ability. Is there something specific you would like to know?
```

名前はアシスタントさんでした。ただ、OpenAIからトレーニングされたという回答は、どこから来ているのか気になる。

日本語でも聞いてみます。

```python
messages = [{"role": "user", "content": "Databricksとは何ですか？詳細に教えてください。"}]
response = model.chat(tokenizer, messages, streaming=False)
print(response)
```

```:出力
Databricksは、クラウドベースの分析サービスと機械学習プラットフォームを提供する企業です。このプラットフォームにより、ユーザーはデータの探索、分析、共有を行うことができます。また、ユーザーはPythonやRなどのプログラミング言語を使用して、自動化されたワークフローを作成し、大規模なデータセットを操作することができます。さらに、Databricksは、クラウドベースのデータレイクと分析サービスを提供しています。
```

的確に回答してくれました。
日本語の文章も破綻がなく、きれいな文章です。

他にも聞いてみます。

```python
messages = [{"role": "user", "content": "まどか☆マギカでは誰が一番かわいい?"}]
response = model.chat(tokenizer, messages, streaming=False)
print(response)
```

```:出力
美樹さやか、巴マミ、佐倉杏子、沙英、藤村かえでの中から選ぶのは難しいですが、個人の好みによって異なります。
```

こちらも、日本語としては綺麗な文章で出力してくれました。
特に指定が無い場合、割と短い文章で回答を生成する傾向があるように思います。

コード生成はどうでしょうか。

```python
messages = [{"role": "user", "content": "ランダムな10個の要素からなるリストを作成してソートするコードをPythonで書いてください。"}]
response = model.chat(tokenizer, messages, streaming=False)
print(response)
```

```:出力
'import random

my_list = []
for i in range(10):
    my_list.append(random.randint(1, 10))

my_list.sort()
print(my_list)'

「import random」というコードを実行し、ランダムな数値を生成するために使用します。
「my_list」という空のリストを作成します。
「for」ループを使用して、10回繰り返し、「random.randint(1, 10)」を使用して1から10までのランダムな数値を生成し、リストに追加します。
リストをソートするために「sort()」メソッドを使用します。
ソートされたリストを出力します。
```

シングルクォーテーションで囲むところが少し独特ですが、説明含めてそれっぽい内容が出力されています。

最後に時事的な質問。

```python
messages = [{"role": "user", "content": "現在の日本の首相は誰？"}]
response = model.chat(tokenizer, messages, streaming=False)
print(response)
```

```:出力
現在の日本の首相は菅義偉です。
```

さすがにLLM単体では正解できませんでしたね。
学習データは少し前のものが中心なのでしょう。

# まとめ

OrionStar社のOrion-14B-ChatをDatabricks上で試してみました。

日本語も、端的に綺麗な文章を作ってくれるため、日系会社のビジネス用途にも使いやすいモデルという印象を持ちました。申請が必要ですが、商用利用も可能ですし。
LLMにおける最近の中心的な用途はRAGだと思うので、RAG用モデルでの利用も試してみたいと思います。

日本語が適切に扱えるモデルがどんどん増えていくのはよいですね！