クラウドMacでMLXのトークナイザーと会話テンプレートを検証する

同じ会話をクラウドMacに移したところ、MLXの回答に突然ロールタグが混ざるようになっても、原因がモデルの重みにあるとは限りません。トークナイザーのバージョン、会話テンプレート、特殊トークンの重複追加によって、実際の入力が変わっている可能性があります。まず「メッセージからトークンへの変換」を固定し、そのうえで生成結果を比較することで、入力のずれとサンプリングのランダム性を切り分けられます。

検証対象となる入力契約を定める

ここで検証するのは推論前の前処理だけです。重みの読み込みや tokens/s の測定は行いません。Pythonの仮想環境と、内容を確認済みのローカルモデルディレクトリを用意してください。ディレクトリには、そのモデルに必要なトークナイザー設定、語彙、会話テンプレートがそろっている必要があります。同規模の別モデルからこれらのファイルを流用しないでください。

1回の検証で記録するのは、元のメッセージ配列、テンプレート適用後のテキスト、完全なトークンID列の3段階です。トークン数を比べるだけでは不十分です。長さが同じ列でも、内容は異なる場合があります。通常の語彙が同じでも、ロールの境界や終了マーカーまで同じとは限りません。

会話テンプレートは表示形式ではなく、モデルの入力プロトコルの一部です。トークン化に成功したというだけでは、変換が完了したことしか分かりません。モデルに正しいロール構造が渡されたことまでは保証されません。

以下のスクリプトを check_prompt.py として保存し、transformers と、そのトークナイザーが必要とする依存パッケージをインストールした環境で実行します。先に python -m pip freeze > requirements.lock.txt を実行して、環境を記録してください。この一覧はバージョンの記録であり、異なるプラットフォームでまったく同じインストール結果になることを保証するものではありません。

ローカルファイルとテスト用メッセージを固定する

スクリプトでは、ネットワーク経由のダウンロードとカスタムのリモートコード実行を明示的に無効にしています。モデルがカスタム実装に依存する場合は、先にその実装をレビューし、別途実行手順を整えてください。エラーを回避するために、その場で実行を許可してはいけません。

サンプルには中国語、改行、山括弧、複数ターンの会話を含め、生成用プレフィックスを付ける場合と付けない場合をそれぞれ確認します。ロール構成は一般的な user と assistant の交互のやり取りです。system ロールやツール呼び出しは、モデルが明示的に対応していることを確認してから、別のサンプルとして追加してください。

import argparse
import json
from pathlib import Path
from transformers import AutoTokenizer

parser = argparse.ArgumentParser()
parser.add_argument("model_dir")
parser.add_argument("baseline")
parser.add_argument("--record", action="store_true")
args = parser.parse_args()

tokenizer = AutoTokenizer.from_pretrained(
    args.model_dir,
    local_files_only=True,
    trust_remote_code=False,
)
cases = [
    [{"role": "user", "content": "只回答数字：2+3\n<note>保留原文</note>"}],
    [
        {"role": "user", "content": "记住代号：实验甲"},
        {"role": "assistant", "content": "已记录。"},
        {"role": "user", "content": "代号是什么？"},
    ],
]
rows = []
for messages in cases:
    for generation in (False, True):
        rendered = tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=generation,
        )
        ids = tokenizer.apply_chat_template(
            messages,
            tokenize=True,
            add_generation_prompt=generation,
        )
        encoded = tokenizer.encode(
            rendered, add_special_tokens=False
        )
        if ids != encoded:
            raise SystemExit("Tokenization paths differ")
        rows.append({
            "messages": messages,
            "generation": generation,
            "rendered": rendered,
            "ids": ids,
        })

path = Path(args.baseline)
if args.record:
    with path.open("x", encoding="utf-8") as stream:
        json.dump(rows, stream, ensure_ascii=False, indent=2)
else:
    expected = json.loads(path.read_text(encoding="utf-8"))
    if rows != expected:
        raise SystemExit("Prompt contract changed")

ベースラインを作成し、自動では承認しない

初回は python check_prompt.py ./model ./prompt-baseline.json --record を実行します。記録時は排他的な新規作成モードを使うため、既存のファイルがあるとエラーになります。これにより、障害調査中に証拠を上書きするのを防ぎます。以降は --record を外し、コマンドが何も出力せず、終了コードがゼロになった場合にのみ合格とします。

JSONを項目ごとに確認してください。ユーザーの内容が欠けていないか、ロールの区切りがモデルの仕様に沿っているか、assistant の過去の応答の後に適切な境界があるか、生成モードで次の応答を開始できる状態になっているかを確認します。生成用プレフィックスの切り替えを無視するテンプレートもあるため、2つの結果が必ず異なることを条件にしてはいけません。

ベースライン、スクリプト、依存パッケージ一覧、使用したトークナイザーファイルのチェックサムをまとめてバージョン管理します。サンプルには合成データだけを使ってください。実際のチャット、アクセストークン、社内文書を、コミット対象のベースラインに含めてはいけません。

実際のMLX呼び出し経路で入力のずれを追う

スクリプトが合格しても、アプリケーションの入口で同じテンプレートとパラメータが使われているかを確認する必要があります。呼び出し層によって、メッセージ配列を受け取るものと、テンプレート適用済みの文字列を受け取るものがあります。その文字列をさらに会話用ラッパーに渡すと、テンプレートが二重に適用される可能性があります。生成関数に入る直前で、合成サンプルの最終的なトークンID列を取得し、ベースラインと要素ごとに比較してください。

特殊トークンの二重追加を防ぐ

テンプレートには開始マーカーや終了マーカーがすでに含まれていることが多いため、スクリプトではテンプレート適用後のテキストをエンコードする際に add_special_tokens=False を指定しています。BOSを手動で追加しないでください。また、デコード結果が「正常に見える」だけで合格と判断してはいけません。特殊トークンを非表示にするデコードオプションによって、入力の重複が見えなくなることがあります。

差分が出たら、まず rendered、次に ids を比較します。テキストが異なる場合は、テンプレートとメッセージの前処理を優先して調べます。テキストが同じでIDだけが異なる場合は、語彙、追加トークン、トークナイザーの実装を確認します。実際の入力が一致したことを確認してから、サンプリングパラメータや重みのバージョンを調べてください。

失敗を再検証できるチェック項目にする

空の内容、末尾の空白、絵文字、長い中国語のサンプルを追加するときは、それぞれ何を検証するのかを明記してください。ランダムな文字列を増やすだけでは不十分です。長い入力については、アプリケーション側でコンテキストの予算も確認する必要があります。このスクリプトは、生成タスクの出力用領域を確保せず、切り詰め方針も検証しません。

引き渡し前に、次の4点を確認します。オフライン実行でファイルがダウンロードされないこと、通常のチェックでベースラインが書き換わらないこと、テンプレートとトークンのどちらに差分があっても処理を停止できること、依存パッケージの更新には人によるレビューが必要であることです。ベースラインを更新する際は、候補を別ファイルに保存して差分を提出し、失敗時の分岐で自動的に再記録しないでください。

レンタル期間中に残すべきものは、ある回答のスクリーンショットではなく、何度でも再実行できる入力の証跡一式です。そうしておけば、別のクラウドMacに移ったときも、まずモデルに渡った内容が変わっていないことを証明してから、回答が異なる理由を検討できます。

よくある質問