チャットテンプレート

チャットの基礎ガイドでは、TextGenerationPipelineを使用してチャット履歴を保存し、チャットモデルからテキストを生成する方法について説明しています。

このガイドは、より高度なユーザーを対象としており、モデルとのチャット時に実際に何が起きているかを理解するための重要な概念に加えて、その基盤となるクラスやメソッドについても解説しています。

チャットモデルを理解するために不可欠な洞察は次の通りです。チャット用に学習されているかどうかに関わらず、すべての因果言語モデル（causal LM）はトークンのシーケンスを継続します。因果言語モデルの学習では、通常、膨大なテキストコーパスを用いた「事前学習（pre-training）」から始まり、「ベース」モデルが作成されます。その後、これらのベースモデルはチャット用に「ファインチューニング」されることが多く、これはメッセージのシーケンスとしてフォーマットされたデータでモデルを学習させることを意味します。ただし、チャットも依然として単なるトークンのシーケンスです！チャットモデルに渡すroleとcontentの辞書のリストは、<|user|>や<|assistant|>、<|end_of_message|>のような制御トークンを用いて、トークンシーケンスに変換されます。これにより、モデルはチャットの構造を認識できるようになります。チャットフォーマットには多くのバリエーションがあり、同じベースモデルからファインチューニングされた場合でも、モデルによって異なるフォーマットや制御トークンが使用されることがあります！

ご安心ください。チャットモデルを使用するために、あらゆるチャット形式を暗記する必要はありません。チャットモデルにはチャットテンプレートが用意されており、チャットがどのようにフォーマットされることを期待しているかを示しています。これらはapply_chat_templateメソッドでアクセスできます。2つの例を見てみましょう。これら2つのモデルは、同じMistral-7Bベースモデルからファインチューニングされています。

Mistral

Zephyr

Mistral-7B-Instruct は [INST] と [/INST] トークンを使用してユーザーメッセージの開始と終了を示し、Zephyr-7B は <|user|> と <|assistant|> トークンを使用して話者の役割を示します。そのため、チャットテンプレートは重要です。制御トークンが誤っていると、これらのモデルの性能は大幅に低下してしまいます。

apply_chat_templateの使用

apply_chat_templateへの入力には、roleキーとcontentキーを持つ辞書のリストとして構造化する必要があります。roleキーは話者を指定し、contentキーにはメッセージが含まれます。一般的なロールは以下の通りです:

ユーザーからのメッセージにはuser
モデルからのメッセージ用のassistant
system は、モデルがどのように動作すべきかを示す指示です（通常、チャットの冒頭に配置されます）。

apply_chat_template はこのリストを受け取り、フォーマットされたシーケンスを返します。シーケンスをトークナイズしたい場合は、tokenize=True を設定してください。

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("HuggingFaceH4/zephyr-7b-beta")
model = AutoModelForCausalLM.from_pretrained("HuggingFaceH4/zephyr-7b-beta", device_map="auto", dtype=torch.bfloat16)

messages = [
    {"role": "system", "content": "You are a friendly chatbot who always responds in the style of a pirate",},
    {"role": "user", "content": "How many helicopters can a human eat in one sitting?"},
 ]
tokenized_chat = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt").to(model.device)
print(tokenizer.decode(tokenized_chat["input_ids"][0]))

<|system|>
You are a friendly chatbot who always responds in the style of a pirate</s>
<|user|>
How many helicopters can a human eat in one sitting?</s>
<|assistant|>

トークン化されたチャットをgenerate()に渡して、応答を生成します。

outputs = model.generate(**tokenized_chat, max_new_tokens=128)
print(tokenizer.decode(outputs[0]))

<|system|>
You are a friendly chatbot who always responds in the style of a pirate</s>
<|user|>
How many helicopters can a human eat in one sitting?</s>
<|assistant|>
Matey, I'm afraid I must inform ye that humans cannot eat helicopters. Helicopters are not food, they are flying machines. Food is meant to be eaten, like a hearty plate o' grog, a savory bowl o' stew, or a delicious loaf o' bread. But helicopters, they be for transportin' and movin' around, not for eatin'. So, I'd say none, me hearties. None at all.

一部のトークナイザーは、<bos> や <eos> などの特殊トークンを追加します。チャットテンプレートには通常、必要な特殊トークンがすべて含まれているため、追加の特殊トークンを加えると誤りや重複が生じ、モデルの性能を損なうことがよくあります。apply_chat_template(tokenize=False) でテキストを整形する場合、後でトークナイズする際はこれらのトークンの重複を避けるため、add_special_tokens=False を設定してください。 apply_chat_template(tokenize=True) を使用する場合、この問題は発生しません。そのため、通常はより安全な選択肢です！

add_generation_prompt

上記の例で、add_generation_prompt引数に気づいたかもしれません。この引数は、チャットの末尾にassistantの応答の開始を示すトークンを追加します。覚えておいてください: チャットの抽象化の奥底には、トークンのシーケンスを続けるだけの言語モデルが存在しています! assistantの応答にいることを示すトークンを含めれば、モデルは正しく応答を書き出しますが、これらのトークンを含めないと、モデルが混乱して、返信する代わりにユーザーのメッセージを続けるような奇妙な動作をする可能性があります。

add_generation_prompt が実際に何をしているのかを理解するために、例を見てみましょう。まず、add_generation_prompt を使わずにチャットをフォーマットします。

tokenized_chat = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)
tokenized_chat

<|im_start|>user
Hi there!<|im_end|>
<|im_start|>assistant
Nice to meet you!<|im_end|>
<|im_start|>user
Can I ask a question?<|im_end|>

では、add_generation_prompt=True を指定して同じチャットをフォーマットしてみましょう。

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

<|im_start|>user
Hi there!<|im_end|>
<|im_start|>assistant
Nice to meet you!<|im_end|>
<|im_start|>user
Can I ask a question?<|im_end|>
<|im_start|>assistant

add_generation_prompt=True の場合、末尾に <|im_start|>assistant が追加され、assistant メッセージの開始を示します。これにより、モデルは次に assistant の応答が続くことを認識できます。

すべてのモデルが生成プロンプトを必要とするわけではなく、Llamaのような一部のモデルでは、assistantの応答の前に特別なトークンがありません。このような場合、add_generation_promptは効果を持ちません。

continue_final_message

continue_final_message パラメータは、チャットの最後のメッセージを新しいメッセージとして開始するのではなく、継続するかどうかを制御します。シーケンス終了トークンを削除することで、モデルが最後のメッセージから生成を継続できるようにします。

これはモデルの応答を「プリフィル」するために有用です。以下の例では、モデルは新しいメッセージを開始するのではなく、JSON文字列を継続するテキストを生成します。応答の開始方法が既知の場合、指示の追従精度を向上させるために非常に有用です。

chat = [
    {"role": "user", "content": "Can you format the answer in JSON?"},
    {"role": "assistant", "content": '{"name": "'},
]

formatted_chat = tokenizer.apply_chat_template(chat, tokenize=True, return_dict=True, continue_final_message=True)
model.generate(**formatted_chat)

add_generation_promptとcontinue_final_messageを同時に使用しないでください。前者は新しいメッセージを開始するトークンを追加し、後者はシーケンス終了トークンを削除します。両方を同時に使用するとエラーが返されます。

content ではなく、フィールド名を文字列として渡すことで、そのフィールドをプリフィルできます。推論モデルは、別の reasoning_content フィールドから推論を読み取ります。content をプリフィルすると、生成が開始される前に推論ブロックが閉じられるため、モデルはブロック内で続行できません。一方、reasoning_content を直接プリフィルすると、ブロックは開いたままになります。

chat = [
    {"role": "user", "content": "Explain 1+1"},
    {"role": "assistant", "reasoning_content": "The user wants a simple addition. ", "content": ""},
]

formatted_chat = tokenizer.apply_chat_template(chat, tokenize=False, continue_final_message="reasoning_content")

名前付きフィールドは最終メッセージに存在する必要があり、チャットテンプレートによって参照される必要があります。いずれかのチェックに失敗すると、エラーが発生します。

TextGenerationPipeline は、新しいメッセージを開始するために、デフォルトで add_generation_prompt を True に設定します。ただし、チャットの最後のメッセージが assistant ロールを持つ場合、そのメッセージはプリフィルであると判断し、continue_final_message=True に切り替えます。これは、ほとんどのモデルが連続する複数の assistant メッセージをサポートしていないためです。この動作を上書きするには、パイプラインに continue_final_message 引数を明示的に渡してください。

モデルの学習

チャットテンプレート付きでモデルを学習させることは、テンプレートがモデルの学習に使われたトークンと一致することを保証する良い方法です。データセットへの前処理ステップとしてチャットテンプレートを適用してください。アシスタントの応答を促すための追加トークンは学習中に役立たないため、add_generation_prompt=Falseを設定してください。

チャットテンプレートを使用してデータセットを前処理する例を以下に示します。

from transformers import AutoTokenizer
from datasets import Dataset

tokenizer = AutoTokenizer.from_pretrained("HuggingFaceH4/zephyr-7b-beta")

chat1 = [
    {"role": "user", "content": "Which is bigger, the moon or the sun?"},
    {"role": "assistant", "content": "The sun."}
]
chat2 = [
    {"role": "user", "content": "Which is bigger, a virus or a bacterium?"},
    {"role": "assistant", "content": "A bacterium."}
]

dataset = Dataset.from_dict({"chat": [chat1, chat2]})
dataset = dataset.map(lambda x: {"formatted_chat": tokenizer.apply_chat_template(x["chat"], tokenize=False, add_generation_prompt=False)})
print(dataset['formatted_chat'][0])

<|user|>
Which is bigger, the moon or the sun?</s>
<|assistant|>
The sun.</s>

このステップの後、formatted_chatカラムを使用して、因果言語モデルのトレーニングレシピに従って続行できます。

Update on GitHub

Transformers

チャットテンプレート

apply_chat_templateの使用

add_generation_prompt

continue_final_message

モデルの学習