From f149c00c9231e13ee7e267c82316d7ec3b87445b Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 20 Dec 2025 07:26:51 +0000 Subject: [PATCH] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- unsloth/models/sentence_transformer.py | 163 +++++++++++++++++-------- 1 file changed, 110 insertions(+), 53 deletions(-) diff --git a/unsloth/models/sentence_transformer.py b/unsloth/models/sentence_transformer.py index 0181fc006e..c328bdb9af 100644 --- a/unsloth/models/sentence_transformer.py +++ b/unsloth/models/sentence_transformer.py @@ -29,7 +29,8 @@ import transformers from packaging.version import Version from transformers import AutoModel, AutoConfig -class FastSentenceTransformer(FastModel): + +class FastSentenceTransformer(FastModel): @staticmethod def _read_pooling_mode(model_name, token): """ @@ -88,7 +89,9 @@ class FastSentenceTransformer(FastModel): return mode except Exception as e: - print(f"\033[1;33mFailed to detect pooling mode, not a sentence-transformers model. You will have to handle pooling/normalization yourself for inference, but training should be fine.\033[0m") + print( + f"\033[1;33mFailed to detect pooling mode, not a sentence-transformers model. You will have to handle pooling/normalization yourself for inference, but training should be fine.\033[0m" + ) return "mean" # should prolly be done upstream instead of this hackfest here @@ -139,9 +142,7 @@ class FastSentenceTransformer(FastModel): def create_custom_forward(module): # bog standard checkpoint def custom_forward(*inputs): - return module( - *inputs, output_attentions = output_attentions - ) + return module(*inputs, output_attentions = output_attentions) return custom_forward @@ -218,8 +219,7 @@ class FastSentenceTransformer(FastModel): output_hidden_states: bool = False, return_dict: bool = False, **kwargs, - ): - + ): position_bias = self.compute_position_bias(hidden_states) all_hidden_states = () if output_hidden_states else None all_attentions = () if output_attentions else None @@ -235,6 +235,7 @@ class FastSentenceTransformer(FastModel): # checkpoint def custom_forward(*inputs): return module(*inputs, output_attentions = output_attentions) + return custom_forward layer_outputs = torch.utils.checkpoint.checkpoint( @@ -283,6 +284,7 @@ class FastSentenceTransformer(FastModel): Patch the forward method of the DistilBertModel to use positional arguments instead of keyword arguments. Transformers 4 version. """ + # based on: # https://github.com/huggingface/transformers/blob/v4.57.3/src/transformers/models/distilbert/modeling_distilbert.py#L666 # original code from here on: @@ -296,21 +298,33 @@ class FastSentenceTransformer(FastModel): output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, ): - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + output_attentions = ( + output_attentions + if output_attentions is not None + else self.config.output_attentions + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict ) - return_dict = return_dict if return_dict is not None else self.config.use_return_dict if input_ids is not None and inputs_embeds is not None: - raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + raise ValueError( + "You cannot specify both input_ids and inputs_embeds at the same time" + ) elif input_ids is not None: self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask) input_shape = input_ids.size() elif inputs_embeds is not None: input_shape = inputs_embeds.size()[:-1] else: - raise ValueError("You have to specify either input_ids or inputs_embeds") + raise ValueError( + "You have to specify either input_ids or inputs_embeds" + ) device = input_ids.device if input_ids is not None else inputs_embeds.device @@ -318,17 +332,29 @@ class FastSentenceTransformer(FastModel): # Prepare head mask if needed head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers) - embeddings = self.embeddings(input_ids, inputs_embeds) # (bs, seq_length, dim) + embeddings = self.embeddings( + input_ids, inputs_embeds + ) # (bs, seq_length, dim) if self.config._attn_implementation == "flash_attention_2": - attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None + attention_mask = ( + attention_mask + if (attention_mask is not None and 0 in attention_mask) + else None + ) else: if attention_mask is None: - attention_mask = torch.ones(input_shape, device=device) # (bs, seq_length) + attention_mask = torch.ones( + input_shape, device = device + ) # (bs, seq_length) - if self.config._attn_implementation == "sdpa" and head_mask_is_none and not output_attentions: + if ( + self.config._attn_implementation == "sdpa" + and head_mask_is_none + and not output_attentions + ): attention_mask = _prepare_4d_attention_mask_for_sdpa( - attention_mask, embeddings.dtype, tgt_len=input_shape[1] + attention_mask, embeddings.dtype, tgt_len = input_shape[1] ) # patch here, change kwargs to positional args: return self.transformer( @@ -352,6 +378,7 @@ class FastSentenceTransformer(FastModel): # https://github.com/huggingface/transformers/blob/v5.0.0rc1/src/transformers/models/distilbert/modeling_distilbert.py#L386 # original code from here on: from transformers.masking_utils import create_bidirectional_mask + def forward( self, input_ids: Optional[torch.Tensor] = None, @@ -361,14 +388,16 @@ class FastSentenceTransformer(FastModel): **kwargs, ): if (input_ids is None) ^ (inputs_embeds is not None): - raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + raise ValueError( + "You must specify exactly one of input_ids or inputs_embeds" + ) embeddings = self.embeddings(input_ids, inputs_embeds, position_ids) attention_mask = create_bidirectional_mask( - config=self.config, - input_embeds=embeddings, - attention_mask=attention_mask, + config = self.config, + input_embeds = embeddings, + attention_mask = attention_mask, ) # patch here: unsloth gradient checkpointing hook needs positional arguments @@ -377,10 +406,11 @@ class FastSentenceTransformer(FastModel): attention_mask, **kwargs, ) + modeling_distilbert.DistilBertModel.forward = forward @staticmethod - def _module_path(model_name, token=None): + def _module_path(model_name, token = None): """ Returns the path to the modules.json file or None """ @@ -390,7 +420,7 @@ class FastSentenceTransformer(FastModel): return path if os.path.exists(path) else None else: try: - return hf_hub_download(model_name, "modules.json", token=token) + return hf_hub_download(model_name, "modules.json", token = token) except: return None except: @@ -406,7 +436,7 @@ class FastSentenceTransformer(FastModel): ): """Helper to create and configure a Transformer module.""" from sentence_transformers.models import Transformer - + transformer_module = Transformer( model_name, max_seq_length = max_seq_length, @@ -416,7 +446,7 @@ class FastSentenceTransformer(FastModel): transformer_module.auto_model = model transformer_module.tokenizer = tokenizer transformer_module.do_lower_case = getattr(tokenizer, "do_lower_case", False) - + # sentence-transformers only passes along known keys to model.forward model_forward_params = list(inspect.signature(model.forward).parameters) transformer_module.model_forward_params = set(model_forward_params) | { @@ -425,23 +455,25 @@ class FastSentenceTransformer(FastModel): "token_type_ids", "inputs_embeds", } - + # determine max_seq_length if not provided if max_seq_length is None: - if hasattr(model, "config") and hasattr(model.config, "max_position_embeddings"): + if hasattr(model, "config") and hasattr( + model.config, "max_position_embeddings" + ): max_seq_length = model.config.max_position_embeddings elif hasattr(tokenizer, "model_max_length"): max_seq_length = tokenizer.model_max_length else: max_seq_length = 512 - + transformer_module.max_seq_length = max_seq_length transformer_module.config_keys = ["max_seq_length", "do_lower_case"] transformer_module.save_in_root = True - + if hasattr(model, "config"): model.config.tokenizer_class = tokenizer.__class__.__name__ - + return transformer_module @staticmethod @@ -456,27 +488,35 @@ class FastSentenceTransformer(FastModel): ) -> tuple[OrderedDict, bool]: """ Load modules from modules.json if available, otherwise fallback to hard-coded modules. - + Returns: tuple[OrderedDict, bool]: (modules, no_modules_json) """ from sentence_transformers.util import import_from_string, load_dir_path from sentence_transformers.models import Pooling, Normalize - + modules = OrderedDict() modules_json_path = FastSentenceTransformer._module_path(model_name, token) - + if modules_json_path: - with open(modules_json_path, encoding="utf8") as f: + with open(modules_json_path, encoding = "utf8") as f: modules_config = json.load(f) for module_config in modules_config: class_ref = module_config["type"] - name = module_config.get("name", str(module_config.get("idx", len(modules)))) + name = module_config.get( + "name", str(module_config.get("idx", len(modules))) + ) if class_ref == "sentence_transformers.models.Transformer": - transformer_module = FastSentenceTransformer._create_transformer_module( - model_name, model, tokenizer, max_seq_length, trust_remote_code + transformer_module = ( + FastSentenceTransformer._create_transformer_module( + model_name, + model, + tokenizer, + max_seq_length, + trust_remote_code, + ) ) modules[name] = transformer_module else: @@ -486,9 +526,13 @@ class FastSentenceTransformer(FastModel): load_path = os.path.join(model_name, module_path) else: try: - load_path = load_dir_path(model_name, module_path, token=token) + load_path = load_dir_path( + model_name, module_path, token = token + ) except Exception as e: - print(f"Unsloth Warning: Could not download module {module_path}: {e}") + print( + f"Unsloth Warning: Could not download module {module_path}: {e}" + ) continue module_class = import_from_string(class_ref) @@ -496,13 +540,17 @@ class FastSentenceTransformer(FastModel): module = module_class.load(load_path) modules[name] = module except Exception as e: - print(f"Unsloth Warning: Failed to load module {name} ({class_ref}): {e}") - + print( + f"Unsloth Warning: Failed to load module {name} ({class_ref}): {e}" + ) + return modules, False - + # fallback if no modules.json (non sentence-transformers models) - print("Unsloth: No modules.json found, falling back to [Transformer, Pooling, Normalize]") - + print( + "Unsloth: No modules.json found, falling back to [Transformer, Pooling, Normalize]" + ) + transformer_module = FastSentenceTransformer._create_transformer_module( model_name, model, tokenizer, max_seq_length, trust_remote_code ) @@ -513,9 +561,11 @@ class FastSentenceTransformer(FastModel): if pooling_mode == "mean": pooling_mode = FastSentenceTransformer._read_pooling_mode(model_name, token) - modules["1"] = Pooling(word_embedding_dimension=hidden_size, pooling_mode=pooling_mode) + modules["1"] = Pooling( + word_embedding_dimension = hidden_size, pooling_mode = pooling_mode + ) modules["2"] = Normalize() - + return modules, True @staticmethod @@ -557,7 +607,9 @@ class FastSentenceTransformer(FastModel): # if for_inference == True, skip Unsloth optimizations to avoid torch compile issues if for_inference: - st_model = SentenceTransformer(model_name, device=device_map, trust_remote_code=trust_remote_code) + st_model = SentenceTransformer( + model_name, device = device_map, trust_remote_code = trust_remote_code + ) return st_model if "auto_model" not in kwargs: @@ -568,7 +620,7 @@ class FastSentenceTransformer(FastModel): if not is_distilbert: try: # this becomes necessary after merging a trained model - config = AutoConfig.from_pretrained(model_name, token=token) + config = AutoConfig.from_pretrained(model_name, token = token) if getattr(config, "model_type", "") == "distilbert": is_distilbert = True except: @@ -599,7 +651,9 @@ class FastSentenceTransformer(FastModel): # check if modules.json exists - if not, force 16-bit training # why? because i have to implement saving myself for these models, and i don't feel like adding dequantization # to the save_pretrained_merged for a model that really should be trained in 16-bit anyway - has_modules_json = FastSentenceTransformer._module_path(model_name, token) is not None + has_modules_json = ( + FastSentenceTransformer._module_path(model_name, token) is not None + ) if not has_modules_json and load_in_4bit: print( @@ -642,6 +696,7 @@ class FastSentenceTransformer(FastModel): # try to load modules, otherwise fallback to old hard-coded modules from sentence_transformers import SentenceTransformer + modules, no_modules = FastSentenceTransformer._load_modules( model_name, token, @@ -669,16 +724,18 @@ class FastSentenceTransformer(FastModel): tokenizer = kwargs.pop("tokenizer", self.tokenizer) if self.no_modules: # fallback for non-sentence-transformers models - print("Unsloth: No modules detected. Using standard merge_and_unload for saving...") + print( + "Unsloth: No modules detected. Using standard merge_and_unload for saving..." + ) safe_kwargs = kwargs.copy() # filter out Unsloth-specific args that are not in huggingface's save_pretrained unsloth_args = [ "save_method", "temporary_location", - "maximum_memory_usage" + "maximum_memory_usage", ] for k in unsloth_args: - safe_kwargs.pop(k, None) + safe_kwargs.pop(k, None) merged_model = self[0].auto_model.merge_and_unload() merged_model.save_pretrained(save_directory, **safe_kwargs)