Fix FastSentenceTransformer loading with newer sentence-transformers (#5259)
* Fix FastSentenceTransformer compatibility with sentence-transformers 5.4 * Support varied Transformer init signatures Detect Transformer.__init__ parameters and build init kwargs accordingly so trust_remote_code and other args are passed using the correct names. Instead of unconditionally using model_args/config_args, the code now inspects the constructor to decide between model_kwargs/config_kwargs vs model_args/config_args and also sets processor_kwargs or tokenizer_args when present. Initializes Transformer with constructed transformer_kwargs (including max_seq_length) to improve compatibility with different Transformer implementations. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Harden SentenceTransformer path and module checks * Scrub .github/workflows for staging push (matches staging base) * Guard auto_model write in FastSentenceTransformer._apply_torch_compile On sentence-transformers >=5.4 Transformer.auto_model is a read-only @property backed by self.model, so a direct assignment raises AttributeError. The two get_peft_model paths already guard the write with isinstance(getattr(type(...), "auto_model", None), property); the auto-compile path missed the same guard, which broke the default trainer path whenever max_steps >= _compile_threshold. * Add tests for FastSentenceTransformer property guards * Tighten FastSentenceTransformer redirect lifecycle tests Drop a duplicate assertion-less case, remove dead AST extraction helper, and trim unused imports. The remaining six tests cover substitution on match, restoration on constructor exception, passthrough for unrelated names, pathlib.Path normalisation, trailing slash handling, and the no-identifier guard. * Sync .github/workflows with upstream author branch * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Avoid sharing trust_remote_code kwargs dict across constructor buckets In FastSentenceTransformer._create_transformer_module, the same trust_remote_code_kwargs dict was being assigned to model_kwargs, config_kwargs, and processor_kwargs (or model_args / config_args / tokenizer_args) on the Transformer constructor. transformers' from_pretrained code paths (configuration_utils, auto_factory, processing_auto, etc.) call kwargs.pop("trust_remote_code", ...) on the dict they receive, which would drain the shared object and silently strip trust_remote_code from the other buckets. Pass an independent copy to each bucket so subsequent buckets and any pass-through auxiliary loads still see trust_remote_code. * Wire do_lower_case and return_dict through Transformer init for ST 5.4 In FastSentenceTransformer._create_transformer_module: - When Transformer.__init__ accepts do_lower_case (ST 5.4+), pass the unsloth tokenizer's do_lower_case as a constructor kwarg. The existing post-init attribute assignment alone is too late: ST 5.4's __init__ uses do_lower_case to install a Lowercase normalizer on tokenizer.backend_tokenizer.normalizer, which is not re-applied if we only set the attribute after construction. The post-init line is preserved untouched for older ST versions. - Add return_dict to the manually completed model_forward_params set so wrapped models with forward(*args, **kwargs) signatures keep ST's forced dict-like output safety net. ST 5.4's own __init__ unions the forward signature with the same set plus return_dict; the previous override silently dropped it. * Preserve flash-attention forward keys when wrapping ST 5.4 Transformer Sentence-transformers 5.4's Transformer.__init__ calls _can_flatten_inputs() during construction, which augments self.model_forward_params with cu_seq_lens_q, cu_seq_lens_k, max_length_q, max_length_k, seq_idx whenever feature-extraction with text modality, the torch backend, flash-attention 2, and varlen flash-attn support are all available. The post-init override of transformer_module.model_forward_params used to replace the attribute outright, silently dropping those keys so ST's preprocess() filter stripped flash-attn kwargs before reaching model.forward. Snapshot the constructor-populated set first, leave the existing overwrite intact for the forward-signature plus tokenizer keys, and union the snapshot back in so flash-attn forwarding keeps working on ST 5.4. For older sentence-transformers releases the attribute is absent and getattr returns an empty set, leaving behavior unchanged. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han <danielhanchen@gmail.com>
This commit is contained in:
Etherll
GitHub
parent
09505fcc6e
commit
680d43a488
2 changed files with 381 additions and 24 deletions
|
|
@ -0,0 +1,218 @@
|
|||
"""FastSentenceTransformer constructor-redirect lifecycle:
|
||||
- AutoModel/AutoProcessor/AutoTokenizer.from_pretrained are restored even
|
||||
when the Transformer constructor raises (try/finally invariant).
|
||||
- The closure that decides whether to substitute the pre-loaded objects
|
||||
(`is_requested_model_name`) handles HF repo IDs, local paths, trailing
|
||||
slashes, pathlib.Path objects, and missing identifiers correctly.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import pathlib
|
||||
import sys
|
||||
import types
|
||||
|
||||
|
||||
class _FakeAuto:
|
||||
def __init__(self, name):
|
||||
self.name = name
|
||||
self.from_pretrained = self._original
|
||||
|
||||
def _original(self, *args, **kwargs):
|
||||
return ("orig", self.name, args, kwargs)
|
||||
|
||||
|
||||
class _RecordingTransformerOk:
|
||||
last_calls = None
|
||||
|
||||
def __init__(self, model_name, **kwargs):
|
||||
from transformers import AutoModel, AutoProcessor, AutoTokenizer
|
||||
|
||||
type(self).last_calls = {
|
||||
"model": AutoModel.from_pretrained(model_name),
|
||||
"processor": AutoProcessor.from_pretrained(model_name),
|
||||
"tokenizer": AutoTokenizer.from_pretrained(model_name),
|
||||
}
|
||||
|
||||
|
||||
class _RaisingTransformer:
|
||||
def __init__(self, *a, **kw):
|
||||
from transformers import AutoModel
|
||||
|
||||
AutoModel.from_pretrained(a[0] if a else kw.get("model_name_or_path"))
|
||||
raise RuntimeError("simulated init failure")
|
||||
|
||||
|
||||
def _build_driver(transformer_class):
|
||||
transformers_mod = types.ModuleType("transformers")
|
||||
transformers_mod.AutoModel = _FakeAuto("AutoModel")
|
||||
transformers_mod.AutoProcessor = _FakeAuto("AutoProcessor")
|
||||
transformers_mod.AutoTokenizer = _FakeAuto("AutoTokenizer")
|
||||
sys.modules["transformers"] = transformers_mod
|
||||
|
||||
st_root = types.ModuleType("sentence_transformers")
|
||||
st_models = types.ModuleType("sentence_transformers.models")
|
||||
st_models.Transformer = transformer_class
|
||||
sys.modules["sentence_transformers"] = st_root
|
||||
sys.modules["sentence_transformers.models"] = st_models
|
||||
|
||||
captured = {"calls": None}
|
||||
|
||||
def driver(model_name, model, tokenizer):
|
||||
from transformers import AutoModel, AutoProcessor, AutoTokenizer
|
||||
from sentence_transformers.models import Transformer
|
||||
|
||||
def is_requested_model_name(args, kwargs):
|
||||
requested = None
|
||||
if args:
|
||||
requested = args[0]
|
||||
else:
|
||||
requested = kwargs.get("pretrained_model_name_or_path")
|
||||
if requested is None:
|
||||
requested = kwargs.get("model_name_or_path")
|
||||
if requested is None:
|
||||
return False
|
||||
try:
|
||||
requested = os.fspath(requested)
|
||||
expected = os.fspath(model_name)
|
||||
except (TypeError, ValueError):
|
||||
return False
|
||||
if requested == expected:
|
||||
return True
|
||||
try:
|
||||
if os.path.exists(requested) or os.path.exists(expected):
|
||||
return os.path.abspath(requested) == os.path.abspath(expected)
|
||||
except (OSError, TypeError, ValueError):
|
||||
pass
|
||||
return False
|
||||
|
||||
original_model = AutoModel.from_pretrained
|
||||
original_processor = AutoProcessor.from_pretrained
|
||||
original_tokenizer = AutoTokenizer.from_pretrained
|
||||
|
||||
def return_existing_model(*a, **kw):
|
||||
return model if is_requested_model_name(a, kw) else original_model(*a, **kw)
|
||||
|
||||
def return_existing_tokenizer(*a, **kw):
|
||||
return (
|
||||
tokenizer
|
||||
if is_requested_model_name(a, kw)
|
||||
else original_tokenizer(*a, **kw)
|
||||
)
|
||||
|
||||
def return_existing_processor(*a, **kw):
|
||||
return (
|
||||
tokenizer
|
||||
if is_requested_model_name(a, kw)
|
||||
else original_processor(*a, **kw)
|
||||
)
|
||||
|
||||
try:
|
||||
AutoModel.from_pretrained = return_existing_model
|
||||
AutoProcessor.from_pretrained = return_existing_processor
|
||||
AutoTokenizer.from_pretrained = return_existing_tokenizer
|
||||
t = Transformer(model_name)
|
||||
captured["calls"] = getattr(type(t), "last_calls", None)
|
||||
return t
|
||||
finally:
|
||||
AutoModel.from_pretrained = original_model
|
||||
AutoProcessor.from_pretrained = original_processor
|
||||
AutoTokenizer.from_pretrained = original_tokenizer
|
||||
|
||||
return driver, transformers_mod, captured
|
||||
|
||||
|
||||
def test_redirect_substitutes_preloaded_objects_on_match():
|
||||
driver, _mod, captured = _build_driver(_RecordingTransformerOk)
|
||||
sentinel_model = object()
|
||||
sentinel_tok = object()
|
||||
driver("sentence-transformers/all-MiniLM-L6-v2", sentinel_model, sentinel_tok)
|
||||
calls = captured["calls"]
|
||||
assert calls["model"] is sentinel_model
|
||||
assert calls["processor"] is sentinel_tok
|
||||
assert calls["tokenizer"] is sentinel_tok
|
||||
|
||||
|
||||
def test_redirect_restored_on_constructor_exception():
|
||||
driver, transformers_mod, _ = _build_driver(_RaisingTransformer)
|
||||
pre_model = transformers_mod.AutoModel.from_pretrained
|
||||
pre_processor = transformers_mod.AutoProcessor.from_pretrained
|
||||
pre_tokenizer = transformers_mod.AutoTokenizer.from_pretrained
|
||||
|
||||
try:
|
||||
driver("model-id", object(), object())
|
||||
except RuntimeError:
|
||||
pass
|
||||
|
||||
assert transformers_mod.AutoModel.from_pretrained is pre_model
|
||||
assert transformers_mod.AutoProcessor.from_pretrained is pre_processor
|
||||
assert transformers_mod.AutoTokenizer.from_pretrained is pre_tokenizer
|
||||
|
||||
|
||||
def test_redirect_passes_through_for_other_model_names():
|
||||
class _OtherNameTransformer:
|
||||
captured = None
|
||||
|
||||
def __init__(self, model_name, **kw):
|
||||
from transformers import AutoModel
|
||||
|
||||
type(self).captured = AutoModel.from_pretrained("some-other/aux-model")
|
||||
|
||||
driver, *_ = _build_driver(_OtherNameTransformer)
|
||||
sentinel = object()
|
||||
driver("primary/model-id", sentinel, object())
|
||||
assert _OtherNameTransformer.captured is not sentinel
|
||||
assert isinstance(_OtherNameTransformer.captured, tuple)
|
||||
assert _OtherNameTransformer.captured[0] == "orig"
|
||||
|
||||
|
||||
def test_is_requested_model_name_handles_pathlib_path(tmp_path):
|
||||
target = tmp_path / "model_dir"
|
||||
target.mkdir()
|
||||
|
||||
class _PathTransformer:
|
||||
last_calls = None
|
||||
|
||||
def __init__(self, model_name, **kw):
|
||||
from transformers import AutoModel
|
||||
|
||||
type(self).last_calls = AutoModel.from_pretrained(pathlib.Path(model_name))
|
||||
|
||||
driver, *_ = _build_driver(_PathTransformer)
|
||||
sentinel_model = object()
|
||||
driver(str(target), sentinel_model, object())
|
||||
assert _PathTransformer.last_calls is sentinel_model
|
||||
|
||||
|
||||
def test_is_requested_model_name_trailing_slash_local_path(tmp_path):
|
||||
target = tmp_path / "model_dir"
|
||||
target.mkdir()
|
||||
|
||||
class _SlashTransformer:
|
||||
last_calls = None
|
||||
|
||||
def __init__(self, model_name, **kw):
|
||||
from transformers import AutoModel
|
||||
|
||||
type(self).last_calls = AutoModel.from_pretrained(str(target) + "/")
|
||||
|
||||
driver, *_ = _build_driver(_SlashTransformer)
|
||||
sentinel_model = object()
|
||||
driver(str(target), sentinel_model, object())
|
||||
assert _SlashTransformer.last_calls is sentinel_model
|
||||
|
||||
|
||||
def test_is_requested_model_name_returns_false_when_no_identifier():
|
||||
captured = {"args": None}
|
||||
|
||||
class _NoNameTransformer:
|
||||
def __init__(self, model_name, **kw):
|
||||
from transformers import AutoModel
|
||||
|
||||
captured["args"] = AutoModel.from_pretrained(some_other_kwarg = "x")
|
||||
|
||||
driver, *_ = _build_driver(_NoNameTransformer)
|
||||
driver("primary/model-id", object(), object())
|
||||
assert isinstance(captured["args"], tuple)
|
||||
assert captured["args"][0] == "orig"
|
||||
Loading…
Add table
Add a link
Reference in a new issue