mirror of
https://github.com/unslothai/unsloth.git
synced 2026-07-09 15:58:41 +00:00
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 <danielhanchen@gmail.com>
This commit is contained in:
Etherl
GitHub
parent
8205d4c081
commit
5e43c623b9
3 changed files with 222 additions and 3 deletions
122
tests/python/test_fast_sentence_transformer_embedding_parity.py
Normal file
122
tests/python/test_fast_sentence_transformer_embedding_parity.py
Normal file
|
|
@ -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]})"
|
||||
)
|
||||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue