Проверка токенизатора и шаблонов чата MLX на облачном Mac

Если после переноса того же диалога на облачный Mac в ответах MLX внезапно появились метки ролей, причина не обязательно в весах модели. Версия токенизатора, шаблон чата или повторно добавленные специальные токены могут изменить фактические входные данные. Сначала зафиксируйте преобразование сообщений в токены и лишь затем сравнивайте результаты генерации: так можно отделить изменения входных данных от случайности семплирования.

Определите контракт входных данных для проверки

В этой практике проверяется только подготовка входных данных перед инференсом: веса не загружаются, скорость в tokens/s не измеряется. Подготовьте виртуальное окружение Python и проверенный локальный каталог модели. В нём должны находиться конфигурация токенизатора, словарь и шаблон чата, необходимые именно этой модели. Не копируйте эти файлы из другой модели сопоставимого размера.

Для каждой проверки сохраняйте три представления: исходный массив сообщений, текст после применения шаблона и полную последовательность ID токенов. Сравнивать только количество токенов недостаточно: две последовательности одинаковой длины могут передавать разное содержание. Совпадение обычного словаря также не гарантирует совпадения границ ролей и маркеров завершения.

Шаблон чата — часть протокола ввода модели, а не способ оформления текста. Успешная токенизация означает лишь, что преобразование выполнено, но не подтверждает, что модель получила правильную структуру ролей.

Сохраните приведённый ниже скрипт в файл 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; подготовлено ли в режиме генерации место для следующего ответа. Некоторые шаблоны игнорируют переключатель префикса генерации, поэтому нельзя требовать, чтобы результаты двух режимов обязательно различались.

Добавьте в систему контроля версий эталон, скрипт, список зависимостей и контрольные суммы использованных файлов токенизатора. Для примеров используйте только синтетические данные. Реальные переписки, токены доступа и внутренние документы не должны попадать в эталон, предназначенный для коммита.

Найдите источник изменений в реальной цепочке вызовов MLX

Даже если скрипт прошёл проверку, нужно убедиться, что на входе приложения используются те же шаблон и параметры. Одни слои вызова принимают массив сообщений, другие — уже сформированную строку. Если повторно передать такую строку обёртке для чата, шаблон может примениться дважды. Непосредственно перед вызовом функции генерации перехватите итоговую последовательность ID токенов для синтетического примера и поэлементно сравните её с эталоном.

Не добавляйте специальные токены повторно

Шаблон часто уже содержит маркеры начала или завершения, поэтому при кодировании сформированного текста скрипт использует add_special_tokens=False. Не добавляйте BOS вручную и не считайте проверку пройденной лишь потому, что декодированный текст «выглядит нормально»: настройки декодирования, скрывающие специальные токены, могут замаскировать их дублирование во входных данных.

При расхождении сначала сравните rendered, затем ids. Если различается текст, в первую очередь проверьте шаблон и предварительную обработку сообщений. Если текст совпадает, а ID различаются, проверьте словарь, добавленные токены и реализацию токенизатора. К параметрам семплирования и версии весов переходите только после подтверждения полного совпадения фактических входных данных.

Превратите сбои в воспроизводимые проверки

Добавляя примеры с пустым содержимым, пробелами в конце, эмодзи и длинным китайским текстом, указывайте, что именно проверяет каждый из них, а не просто накапливайте случайные строки. Для длинного ввода также проверяйте бюджет контекста на уровне приложения: этот скрипт не резервирует место под ответ и не проверяет стратегию усечения.

Перед передачей результата убедитесь в четырёх вещах: при автономном запуске файлы не скачиваются; обычная проверка не перезаписывает эталон; расхождения как в тексте после применения шаблона, так и в токенах останавливают выполнение; обновление зависимостей требует ручной проверки. Для обновления эталона сохраните кандидат в отдельный файл и отправьте различия на проверку. Не запускайте автоматическую перезапись эталона в ветке обработки сбоя.

На время аренды сохраняйте не скриншот отдельного ответа, а набор входных данных и проверок, которые можно воспроизвести. Тогда при переходе на другой облачный Mac вы сможете сначала доказать, что модель получила те же данные, и лишь затем разбираться, почему изменился ответ.

Часто задаваемые вопросы