From f56780885ac57a259f8c1c25a5e4b7cd7bec7f9e Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 28 Jun 2026 02:36:22 +0000 Subject: [PATCH] Fix model load crash on malformed extra_special_tokens A tokenizer_config.json with extra_special_tokens as a JSON array instead of an object makes transformers raise "'list' object has no attribute 'keys'" in SpecialTokensMixin._set_model_specific_special_tokens, before the model can load (seen on unsloth/Qwen3.6-35B-A3B-MLX-8bit on Apple Silicon). Add install_extra_special_tokens_compat() to coerce a non-dict value to {} during tokenizer init, and call it on the MLX and transformers load paths before from_pretrained. Idempotent, version guarded, and a no-op on well-formed configs. --- studio/backend/core/inference/inference.py | 7 ++ .../backend/core/inference/mlx_inference.py | 6 + studio/backend/tests/test_tokenizer_compat.py | 103 ++++++++++++++++++ studio/backend/utils/tokenizer_compat.py | 55 ++++++++++ 4 files changed, 171 insertions(+) create mode 100644 studio/backend/tests/test_tokenizer_compat.py create mode 100644 studio/backend/utils/tokenizer_compat.py diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py index 4dca4db768..5120fcbc84 100644 --- a/studio/backend/core/inference/inference.py +++ b/studio/backend/core/inference/inference.py @@ -240,6 +240,13 @@ class InferenceBackend: return False self.loading_models.add(model_name) + + # Tolerate malformed (list) extra_special_tokens so transformers does + # not raise "'list' object has no attribute 'keys'" during load. + from utils.tokenizer_compat import install_extra_special_tokens_compat + + install_extra_special_tokens_compat() + device_map = get_device_map(gpu_ids) logger.info( f"Using device_map='{device_map}' ({get_visible_gpu_count()} GPU(s) visible)" diff --git a/studio/backend/core/inference/mlx_inference.py b/studio/backend/core/inference/mlx_inference.py index 5c7799152f..0c0b8043eb 100644 --- a/studio/backend/core/inference/mlx_inference.py +++ b/studio/backend/core/inference/mlx_inference.py @@ -143,6 +143,12 @@ class MLXInferenceBackend: "(unsloth_zoo.mlx.loader). Reinstall via install.sh on Apple Silicon." ) from e + # Tolerate malformed (list) extra_special_tokens before mlx-lm builds the + # tokenizer; else transformers raises "'list' has no attribute 'keys'". + from utils.tokenizer_compat import install_extra_special_tokens_compat + + install_extra_special_tokens_compat() + model, tokenizer_or_processor = FastMLXModel.from_pretrained( model_name, max_seq_length = max_seq_length, diff --git a/studio/backend/tests/test_tokenizer_compat.py b/studio/backend/tests/test_tokenizer_compat.py new file mode 100644 index 0000000000..e08e1ffcd2 --- /dev/null +++ b/studio/backend/tests/test_tokenizer_compat.py @@ -0,0 +1,103 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Regression tests for utils.tokenizer_compat. + +Covers the crash reported on unsloth/Qwen3.6-35B-A3B-MLX-8bit: a model whose +tokenizer_config.json ships extra_special_tokens as a JSON array instead of an +object makes transformers raise "'list' object has no attribute 'keys'" during +tokenizer init. install_extra_special_tokens_compat() coerces a non-dict value +to {} so the load succeeds. +""" + +import logging +import sys +import types +from pathlib import Path + +import pytest + +# Self-contained: put backend root on sys.path and stub the structlog-backed +# ``loggers`` module so the unit under test imports without structlog installed +# (mirrors tests/test_transformers_version.py). +_backend_root = Path(__file__).resolve().parent.parent +if str(_backend_root) not in sys.path: + sys.path.insert(0, str(_backend_root)) + +_loggers_stub = types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: logging.getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) + +import utils.tokenizer_compat as tc +from utils.tokenizer_compat import install_extra_special_tokens_compat + +tub = pytest.importorskip("transformers.tokenization_utils_base") + + +def _make_mixin(): + """A real SpecialTokensMixin instance (its __init__ sets _special_tokens_map).""" + + class _Dummy(tub.SpecialTokensMixin): + pass + + return _Dummy() + + +class _RecordingLogger: + def __init__(self): + self.warnings = [] + + def warning(self, msg, *args, **kwargs): + self.warnings.append(msg % args if args else msg) + + +@pytest.fixture +def recording_logger(monkeypatch): + rec = _RecordingLogger() + monkeypatch.setattr(tc, "logger", rec) + return rec + + +def test_method_exists_and_install_is_idempotent(): + # The method we patch must exist in the installed transformers; otherwise the + # bug shape (and this shim) would not apply. + assert hasattr(tub.SpecialTokensMixin, "_set_model_specific_special_tokens") + assert install_extra_special_tokens_compat() is True + first = tub.SpecialTokensMixin._set_model_specific_special_tokens + # A second install must not re-wrap. + assert install_extra_special_tokens_compat() is True + assert tub.SpecialTokensMixin._set_model_specific_special_tokens is first + assert hasattr(first, "__wrapped__") # original kept reachable + + +def test_original_method_crashes_on_list(): + install_extra_special_tokens_compat() + original = tub.SpecialTokensMixin._set_model_specific_special_tokens.__wrapped__ + d = _make_mixin() + with pytest.raises(AttributeError): + original(d, []) # the exact "'list' object has no attribute 'keys'" shape + + +def test_patched_coerces_empty_list(recording_logger): + install_extra_special_tokens_compat() + d = _make_mixin() + d._set_model_specific_special_tokens([]) # must not raise + assert d.extra_special_tokens == {} + assert len(recording_logger.warnings) == 1 + + +def test_patched_coerces_non_empty_list(recording_logger): + install_extra_special_tokens_compat() + d = _make_mixin() + d._set_model_specific_special_tokens([""]) # must not raise + assert d.extra_special_tokens == {} + assert len(recording_logger.warnings) == 1 + + +def test_patched_preserves_valid_dict(recording_logger): + install_extra_special_tokens_compat() + d = _make_mixin() + d._set_model_specific_special_tokens({"img_token": ""}) + assert "img_token" in d.SPECIAL_TOKENS_ATTRIBUTES + assert d._special_tokens_map.get("img_token") == "" + assert recording_logger.warnings == [] # no coercion for a well-formed dict diff --git a/studio/backend/utils/tokenizer_compat.py b/studio/backend/utils/tokenizer_compat.py new file mode 100644 index 0000000000..3583db7b33 --- /dev/null +++ b/studio/backend/utils/tokenizer_compat.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: AGPL-3.0-only +"""Tokenizer load-time compatibility shims. + +Some converted/quantized checkpoints ship tokenizer_config.json with +``extra_special_tokens`` as a JSON array (``[]``) instead of an object (``{}``); +transformers then runs ``list(special_tokens.keys())`` and raises +"'list' object has no attribute 'keys'". We coerce non-dict values to ``{}``. +""" + +from loggers import get_logger + +logger = get_logger(__name__) + +_PATCH_FLAG = "_unsloth_extra_special_tokens_compat" + + +def install_extra_special_tokens_compat() -> bool: + """Coerce a non-dict ``extra_special_tokens`` to ``{}`` during tokenizer init. + + Wraps ``SpecialTokensMixin._set_model_specific_special_tokens`` so a model whose + tokenizer_config.json has a malformed (array) ``extra_special_tokens`` loads + instead of raising "'list' object has no attribute 'keys'". Idempotent and + cheap; safe to call before every load. Returns True when active, False when the + method is absent (such builds lack the bug). + """ + try: + import transformers.tokenization_utils_base as tub + except Exception: + return False + + mixin = getattr(tub, "SpecialTokensMixin", None) + orig = getattr(mixin, "_set_model_specific_special_tokens", None) + if mixin is None or orig is None: + return False + if getattr(mixin, _PATCH_FLAG, False): + return True + + def _patched(self, special_tokens): + if not isinstance(special_tokens, dict): + logger.warning( + "Coercing malformed extra_special_tokens (%s) to {}; " + "tokenizer_config.json should use an object, not an array.", + type(special_tokens).__name__, + ) + special_tokens = {} + try: + self.extra_special_tokens = {} + except Exception: + pass + return orig(self, special_tokens) + + _patched.__wrapped__ = orig # keep original reachable + mixin._set_model_specific_special_tokens = _patched + setattr(mixin, _PATCH_FLAG, True) + return True