Hugging FaceのChat Template機能を使って、OSSな大規模言語モデルをOpenAIのモデルっぽく使ってみた。

2024.02.20

こんちには。

データアナリティクス事業本部 機械学習チームの中村(nokomoro3)です。

みなさん、OSSの大規模言語モデル(LLM)使っていますか?

多くのOSSの大規模言語モデルはHugging Faceで公開されていることは皆さんご存知かと思いますが、少し前にこちらにChat Templateという機能が追加されていたようです。

Chat Templateを使うことで、OpenAIのLLMのようなフォーマットで統一的にLLMとチャットをすることが可能となります。

以下の記事を見てその存在に気がつきました。

こちらの記事はText Generation Inference (TGI)やInference Endpointsが、OpenAIのLLMと互換性のあるMessage APIを使える旨の紹介でしたが、 少し情報を探してみると、chat_templateが以下のようにtokenizer.jsonに記載されていれば、Chat Templateを使うことが可能となるようですので、今回試してみることにしました。

使用するモデル

PLaMo-13Bの以下のモデルを使用してみます。

PLaMoは、Preferred Networksが開発し2023年9月に公開された、日英2言語を合わせた能力で世界トップレベルの高い性能を示す大規模言語モデル(LLM)です。

PLaMo-13Bは公開データセットのみで学習され、商用利用も可能なApache License 2.0で公開されています。

このPLaMoのplamo-13b-instructは、以下のようにchat_templateが記載されているため、今回試してみるのに適していそうです。

試してみた

使用環境

実行環境はGoogle Colaboratoryを使いました。

バージョンですがPythonは3.10.12で、PyTorchは2.1.0+cu121となっています。

!python --version
!python -c 'import torch; print(torch.__version__)'

# Python 3.10.12
# 2.1.0+cu121

GPUはTesla T4(15GB GPU RAM)を使用しました。

なおTesla T4だとテキスト生成速度はかなり時間が掛かりますので、気になる場合はV100の方を使用することも検討されてください。

!nvidia-smi -L
!nvidia-smi -q -d MEMORY
!nvcc --version

# GPU 0: Tesla T4 (UUID: GPU-b63338df-a085-4d97-2bcb-ce3ea9cc343d)
#
# ==============NVSMI LOG==============
#
# Timestamp                                 : Mon Feb 19 14:48:08 2024
# Driver Version                            : 535.104.05
# CUDA Version                              : 12.2
#
# Attached GPUs                             : 1
# GPU 00000000:00:04.0
#     FB Memory Usage
#         Total                             : 15360 MiB
#         Reserved                          : 257 MiB
#         Used                              : 0 MiB
#         Free                              : 15101 MiB
#     BAR1 Memory Usage
#         Total                             : 256 MiB
#         Used                              : 2 MiB
#         Free                              : 254 MiB
#     Conf Compute Protected Memory Usage
#         Total                             : 0 MiB
#         Used                              : 0 MiB
#         Free                              : 0 MiB
#
# nvcc: NVIDIA (R) Cuda compiler driver
# Copyright (c) 2005-2023 NVIDIA Corporation
# Built on Tue_Aug_15_22:02:13_PDT_2023
# Cuda compilation tools, release 12.2, V12.2.140
# Build cuda_12.2.r12.2/compiler.33191640_0

システムRAMは通常メモリの12GBとしました。

!free -h --si | awk  '/Mem:/{print $2}'

# 12G

Pythonのライブラリとして、以下を追加でインストールしました。(Tesla T4ではbitsandbytesがうまく使えなかったのですが、使う時のために)

!pip install accelerate bitsandbytes

バージョンの確認です。

!pip freeze | grep -e "numpy" -e "sentencepiece" -e "torch" -e "transformers" -e "accelerate" -e "bitsandbytes"

# accelerate==0.27.2
# bitsandbytes==0.42.0
# numpy==1.25.2
# sentencepiece==0.1.99
# torch @ https://download.pytorch.org/whl/cu121/torch-2.1.0%2Bcu121-cp310-cp310-linux_x86_64.whl#sha256=0d4e8c52a1fcf5ed6cfc256d9a370fcf4360958fc79d0b08a51d55e70914df46
# torchaudio @ https://download.pytorch.org/whl/cu121/torchaudio-2.1.0%2Bcu121-cp310-cp310-linux_x86_64.whl#sha256=676bda4042734eda99bc59b2d7f761f345d3cde0cad492ad34e3aefde688c6d8
# torchdata==0.7.0
# torchsummary==1.5.1
# torchtext==0.16.0
# torchvision @ https://download.pytorch.org/whl/cu121/torchvision-0.16.0%2Bcu121-cp310-cp310-linux_x86_64.whl#sha256=e76e78d0ad43636c9884b3084ffaea8a8b61f21129fbfa456a5fe734f0affea9
# transformers==4.35.2

コード

まずはモデルをpipelineとしてHugging Faceからロードします。

import torch
from transformers import pipeline

pipe = pipeline("text-generation", model="pfnet/plamo-13b-instruct",
    torch_dtype=torch.float16, device_map="auto", trust_remote_code=True)

bitsandbytesを使用する場合は、Tesla T4からV100に変更して、model_kwargs={"load_in_8bit": True}などを指定してみてください。

次にOpenAI APIを使う時のようなmessagesを定義し、apply_chat_templateを呼び出すことで、chat_templateを使うことができます。

messages = [
    {
        "role": "system",
        "content": "あなたは献立を考えるシェフです。",
    },
    {"role": "user", "content": "おすすめの肉料理を5つ教えてください。"},
]

prompt = pipe.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

apply_chat_templateの詳細は以下をご確認ください。

あとはpipelineにデータを入力して結果を得ます。

outputs = pipe(prompt, max_new_tokens=256, do_sample=True, temperature=0.01)

結果は以下のように入力プロンプト込みの結果となっています。

print(outputs[0]["generated_text"])
以下はタスクを説明する指示で、文脈を説明した入力とペアになっています。要求を適切に補完するよう応答を書いてください。

### 指示:
あなたは献立を考えるシェフです。

### 入力:
おすすめの肉料理を5つ教えてください。

### 応答:
1.ローストビーフ
2.ローストチキン
3.ポークチョップ
4.ビーフシチュー
5.ローストポーク

応答だけを得るためには、以下のようにすれば良さそうです(単純なPythonの処理です)。

print(outputs[0]["generated_text"].lstrip(prompt))
1.ローストビーフ
2.ローストチキン
3.ポークチョップ
4.ビーフシチュー
5.ローストポーク

トラブルシューティング

はじめはpipelineのインスタンス化時の引数をtorch_dtype=torch.bfloat16としており、以下のエラーが出ていました。

Error: cutlassF: no kernel found to launch #327

ジャストな情報を見つけられなかったのですが、以下にあるようにbfloat16をサポートしていないGPUで実行しようとすると発生するようでしたので、今回はtorch_dtype=torch.float16として実行してみました。

まとめ

いかがでしたでしょうか。

Chat Templateが定義されているLLMであれば、簡単にOpenAIっぽく使うことができることが分かりました。これは便利ですね!!

様々なモデルを使う際にも統一的に使用できるため、比較検討する際にも便利だなと感じました。

本記事がOSSなLLMを使いたいと思われている方の参考になれば幸いです。

参考