ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

[pre-commit.ci] auto fixes from pre-commit.com hooks

Browse source 
for more information, see https://pre-commit.ci
This commit is contained in:
pre-commit-ci[bot] 2025-12-20 07:26:51 +00:00
commit f149c00c92
1 changed files with 110 additions and 53 deletions
Download patch file Download diff file Expand all files Collapse all files

163
unsloth/models/sentence_transformer.py
View file

@ -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)