From 5e43c623b98affc23efbf9dbe71061de7c1706a2 Mon Sep 17 00:00:00 2001 From: Etherl <61019402+Etherll@users.noreply.github.com> Date: Thu, 9 Jul 2026 11:46:22 +0300 Subject: [PATCH] Fix FastSentenceTransformer Qwen embedding preprocessing (#6939) * Fix FastSentenceTransformer Qwen embedding preprocessing * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Document Transformer.load embedding modality fix for #6881 * Harden #6881 fix and add forwards/backwards-compatible regression tests * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fall back to Transformer constructor on legacy sentence-transformers without Hub-capable load * Mirror legacy sentence-transformers fallback in embedding-parity tripwire test * Tighten #6881 comments and docstrings * Skip embedding-parity test on CPU-only runners since FastSentenceTransformer requires CUDA * Honor the transformer module's saved subfolder when loading modules.json records a path for the Transformer module (root for decoder embedders like Qwen3-Embedding, 0_Transformer for the classic layout). Pooling/Normalize already load from their saved path; thread the same path into Transformer.load as subfolder so config and tokenizer resolve like stock ST. stays a no-op, so single-module models are unchanged. * Make embedding-parity test bf16-aware fp16 overflows to NaN on bf16-native embedders such as EmbeddingGemma (Gemma3), producing a false parity failure. Prefer bf16 when the GPU supports it so the tripwire can guard the full documented embedding matrix (Qwen3-Embedding, EmbeddingGemma, BGE-M3, all-MiniLM, GTE-ModernBERT), not just fp16-safe models. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: danielhanchen --- ...t_sentence_transformer_embedding_parity.py | 122 ++++++++++++++++++ ...st_sentence_transformers_pinned_symbols.py | 38 ++++++ unsloth/models/sentence_transformer.py | 65 +++++++++- 3 files changed, 222 insertions(+), 3 deletions(-) create mode 100644 tests/python/test_fast_sentence_transformer_embedding_parity.py diff --git a/tests/python/test_fast_sentence_transformer_embedding_parity.py b/tests/python/test_fast_sentence_transformer_embedding_parity.py new file mode 100644 index 000000000..252d4486a --- /dev/null +++ b/tests/python/test_fast_sentence_transformer_embedding_parity.py @@ -0,0 +1,122 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. +"""Regression guard for issue #6881: FastSentenceTransformer must preprocess text +like a stock SentenceTransformer for decoder embedding models. ST 5.x infers a +"message" modality for chat-template models (e.g. Qwen/Qwen3-Embedding), so building +via `Transformer(model_name, ...)` chat-wraps inputs and degrades embeddings; +`_create_transformer_module` uses `Transformer.load(...)` instead. + +Layers: test_transformer_load_signature_supports_unsloth_kwargs (fast, runs when ST +is importable) and test_fast_sentence_transformer_matches_stock_st (end-to-end parity, +opt-in via UNSLOTH_EMBEDDING_PARITY_MODEL so default CI is unaffected). +""" + +from __future__ import annotations + +import inspect +import os + +import pytest + + +def test_transformer_load_signature_supports_unsloth_kwargs(): + """Forwards-compat tripwire: a Hub-capable Transformer.load must accept the kwargs + the #6881 fix passes. Legacy ST 3.x/4.x expose load(input_path); the code falls back + to Transformer(...) there, so mirror that gate and skip.""" + models = pytest.importorskip("sentence_transformers.models") + load = getattr(models.Transformer, "load", None) + assert callable(load), ( + "sentence_transformers Transformer.load is missing; the #6881 fix in " + "unsloth.models.sentence_transformer._create_transformer_module depends on it." + ) + params = inspect.signature(load).parameters + accepts_var_kw = any(p.kind is inspect.Parameter.VAR_KEYWORD for p in params.values()) + # Mirror _create_transformer_module's hub_capable gate. + hub_capable = accepts_var_kw or any(k in params for k in ("token", "cache_folder", "revision")) + if not hub_capable: + pytest.skip( + "legacy Transformer.load(input_path); production path falls back to Transformer(...)" + ) + unsupported = [ + k + for k in ("token", "cache_folder", "revision", "trust_remote_code") + if not (accepts_var_kw or k in params) + ] + assert not unsupported, ( + f"installed sentence_transformers Transformer.load no longer accepts {unsupported} " + f"and has no **kwargs; update _create_transformer_module (#6881) before it silently " + f"falls back to Transformer(...)." + ) + + +def _probe_texts(): + return [ + "roasted chickpeas in 20 kg bags", + "The capital of France is Paris.", + "A fast brown fox jumps over the lazy dog.", + "recette de tarte aux pommes traditionnelle", + ] + + +def test_fast_sentence_transformer_matches_stock_st(): + """End-to-end: FastSentenceTransformer embeddings and tokenization must match a + stock SentenceTransformer load of the same checkpoint. Opt-in (needs a model) and + GPU-only (FastSentenceTransformer requires CUDA), so it skips on CPU-only runners.""" + model_id = os.environ.get("UNSLOTH_EMBEDDING_PARITY_MODEL") + if not model_id: + pytest.skip( + "set UNSLOTH_EMBEDDING_PARITY_MODEL to a chat-template embedding model " + "(HF id or local path) to run the #6881 parity test" + ) + + torch = pytest.importorskip("torch") + if not torch.cuda.is_available(): + pytest.skip("FastSentenceTransformer requires CUDA; skipping on CPU-only runner") + np = pytest.importorskip("numpy") + pytest.importorskip("sentence_transformers") + from sentence_transformers import SentenceTransformer + + device = "cuda" + # Prefer bf16 when the GPU supports it: fp16 overflows to NaN on bf16-native + # embedders such as EmbeddingGemma (Gemma3), which would mask real parity. + dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16 + texts = _probe_texts() + max_seq_length = 256 + + # Control FIRST, before importing unsloth, so its global import patches never + # touch the stock reference (mirrors the issue's "restart runtime" repro). + ctrl = SentenceTransformer(model_id, device = device, model_kwargs = {"torch_dtype": dtype}) + ctrl.max_seq_length = max_seq_length + ctrl_ids = ctrl.tokenize([texts[0]])["input_ids"][0].tolist() + ctrl_emb = np.asarray( + ctrl.encode(texts, normalize_embeddings = True, batch_size = 8), dtype = np.float32 + ) + + import unsloth # noqa: F401 + from unsloth import FastSentenceTransformer + + fast = FastSentenceTransformer.from_pretrained( + model_id, + max_seq_length = max_seq_length, + dtype = dtype, + load_in_4bit = False, + load_in_16bit = True, + ) + fast_ids = fast.tokenize([texts[0]])["input_ids"][0].tolist() + fast_emb = np.asarray( + fast.encode(texts, normalize_embeddings = True, batch_size = 8), dtype = np.float32 + ) + + # Identical tokenization = no chat-template wrapping slipped in (the #6881 defect). + assert fast_ids == ctrl_ids, ( + f"tokenization diverged (chat-template wrapping regressed?):\n" + f" stock: {ctrl_ids}\n fast: {fast_ids}" + ) + + cos = (ctrl_emb * fast_emb).sum(1) / ( + np.linalg.norm(ctrl_emb, axis = 1) * np.linalg.norm(fast_emb, axis = 1) + ) + assert float(cos.min()) > 0.99, ( + f"embedding parity regressed: min cosine {float(cos.min()):.5f} <= 0.99 " + f"(per-text {[round(float(c), 5) for c in cos]})" + ) diff --git a/tests/version_compat/test_sentence_transformers_pinned_symbols.py b/tests/version_compat/test_sentence_transformers_pinned_symbols.py index c0c35b9d5..d7f9a5481 100644 --- a/tests/version_compat/test_sentence_transformers_pinned_symbols.py +++ b/tests/version_compat/test_sentence_transformers_pinned_symbols.py @@ -19,6 +19,8 @@ ST_TAGS = [ "v5.2.3", "v5.3.0", "v5.4.1", + "v5.5.1", + "v5.6.0", "master", ] @@ -120,6 +122,42 @@ def test_st_transformer_base_class_either_path(tag: str): ) +# Transformer.load classmethod: unsloth builds saved-ST modules through it (#6881). +@pytest.mark.parametrize("tag", ST_TAGS) +def test_st_transformer_load_accepts_unsloth_kwargs(tag: str): + """unsloth builds saved ST models via Transformer.load(...) so the saved + modality_config is honored (#6881). If .load stops accepting the hub kwargs it + passes (and has no **kwargs), update the fix before it silently regresses. Not + locating .load is a SKIP (may be inherited); the live test guards the install.""" + candidates = [ + "sentence_transformers/models/Transformer.py", + "sentence_transformers/models/transformer.py", + "sentence_transformers/base/modules/transformer.py", + "sentence_transformers/base/modules/module.py", + ] + for p in candidates: + src = fetch_text("UKPLab/sentence-transformers", tag, p) + if src is None or not has_def(src, "load", "func"): + continue + m = re.search(r"def\s+load\s*\((.*?)\)\s*(?:->[^:]*)?:", src, re.S) + if m is None: + continue + sig = m.group(1) + accepts_var_kw = "**" in sig + missing = [ + kw + for kw in ("token", "cache_folder", "revision", "trust_remote_code") + if not (accepts_var_kw or re.search(rf"\b{re.escape(kw)}\b", sig)) + ] + assert not missing, ( + f"{tag}: Transformer.load in {p} no longer accepts {missing} and has no " + f"**kwargs; update unsloth.models.sentence_transformer._create_transformer_module " + f"(#6881) before it silently falls back to Transformer(...)." + ) + return + pytest.skip(f"{tag}: Transformer.load not locatable in {candidates} (may be inherited)") + + # sentence_transformers.util: import_from_string + load_dir_path helpers unsloth calls. @pytest.mark.parametrize("tag", ST_TAGS) def test_st_util_helpers(tag: str): diff --git a/unsloth/models/sentence_transformer.py b/unsloth/models/sentence_transformer.py index c1172faa9..4a1a555bc 100644 --- a/unsloth/models/sentence_transformer.py +++ b/unsloth/models/sentence_transformer.py @@ -990,7 +990,17 @@ class FastSentenceTransformer(FastModel): return None @staticmethod - def _create_transformer_module(model_name, model, tokenizer, max_seq_length, trust_remote_code): + def _create_transformer_module( + model_name, + model, + tokenizer, + max_seq_length, + trust_remote_code, + token = None, + cache_dir = None, + revision = None, + module_subfolder = "", + ): """Helper to create and configure a Transformer module.""" from sentence_transformers.models import Transformer @@ -1077,7 +1087,45 @@ class FastSentenceTransformer(FastModel): elif "tokenizer_args" in transformer_init_params: transformer_kwargs["tokenizer_args"] = trust_remote_code_kwargs.copy() - transformer_module = Transformer(model_name, **transformer_kwargs) + # Build via Transformer.load so the saved modality_config is honored: plain + # Transformer(...) makes ST 5.x infer a "message" modality for chat-template + # models (e.g. Qwen3-Embedding), chat-wrapping inputs and degrading embeddings + # (#6881). Only use .load when it resolves a Hub id (accepts the kwargs or + # **kwargs); legacy ST 3.x/4.x load(input_path) is local-only with no modality + # bug, so fall back to the constructor. + transformer_module = None + transformer_load = getattr(Transformer, "load", None) + has_modules_json = ( + FastSentenceTransformer._module_path( + model_name, token, cache_dir = cache_dir, revision = revision + ) + is not None + ) + if callable(transformer_load) and has_modules_json: + load_params = inspect.signature(transformer_load).parameters + accepts_var_kw = any( + p.kind is inspect.Parameter.VAR_KEYWORD for p in load_params.values() + ) + hub_capable = accepts_var_kw or any( + key in load_params for key in ("token", "cache_folder", "revision") + ) + if hub_capable: + load_kwargs = { + "token": token, + "cache_folder": cache_dir, + "revision": revision, + "trust_remote_code": trust_remote_code, + **transformer_kwargs, + } + # Resolve config/tokenizer from the module's saved subfolder + # (modules.json "path"), like stock ST; "" (root) is a no-op. + if module_subfolder: + load_kwargs["subfolder"] = module_subfolder + if not accepts_var_kw: + load_kwargs = {k: v for k, v in load_kwargs.items() if k in load_params} + transformer_module = Transformer.load(model_name, **load_kwargs) + if transformer_module is None: + transformer_module = Transformer(model_name, **transformer_kwargs) finally: # Restore original Auto* loading immediately AutoModel.from_pretrained = original_model_from_pretrained @@ -1191,6 +1239,10 @@ class FastSentenceTransformer(FastModel): tokenizer, max_seq_length, trust_remote_code, + token, + cache_dir, + revision, + module_subfolder = module_config.get("path") or "", ) modules[name] = transformer_module else: @@ -1226,7 +1278,14 @@ class FastSentenceTransformer(FastModel): ) transformer_module = FastSentenceTransformer._create_transformer_module( - model_name, model, tokenizer, max_seq_length, trust_remote_code + model_name, + model, + tokenizer, + max_seq_length, + trust_remote_code, + token, + cache_dir, + revision, ) modules["0"] = transformer_module