diff --git a/unsloth/models/__init__.py b/unsloth/models/__init__.py index d7b2393c89..138f309032 100644 --- a/unsloth/models/__init__.py +++ b/unsloth/models/__init__.py @@ -19,6 +19,7 @@ from .qwen2 import FastQwen2Model from .qwen3 import FastQwen3Model from .qwen3_moe import FastQwen3MoeModel from .granite import FastGraniteModel +from .sentence_transformer import FastSentenceTransformer try: from .falcon_h1 import FastFalconH1Model diff --git a/unsloth/models/sentence_transformer.py b/unsloth/models/sentence_transformer.py new file mode 100644 index 0000000000..6bcf410866 --- /dev/null +++ b/unsloth/models/sentence_transformer.py @@ -0,0 +1,293 @@ +# Copyright 2025 electroglyph. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .loader import FastModel + + +class FastSentenceTransformer(FastModel): + @staticmethod + def from_pretrained( + model_name, + max_seq_length=None, + dtype=None, + load_in_4bit=True, + load_in_8bit=False, + load_in_16bit=False, + full_finetuning=False, + token=None, + device_map="sequential", + rope_scaling=None, + fix_tokenizer=True, + trust_remote_code=False, + use_gradient_checkpointing="unsloth", + resize_model_vocab=None, + revision=None, + use_exact_model_name=False, + offload_embedding=False, + random_state=3407, + max_lora_rank=64, + disable_log_stats=True, + qat_scheme=None, + load_in_fp8=False, + unsloth_tiled_mlp=False, + pooling_mode="mean", + **kwargs, + ): + try: + import sentence_transformers + from sentence_transformers import SentenceTransformer + from sentence_transformers.models import Transformer, Pooling, Normalize + from transformers import AutoModel + except ImportError: + raise ImportError( + "Unsloth: To use `FastSentenceTransformer`, you must install `sentence-transformers`.\n" + "Run `pip install sentence-transformers` to install it." + ) + + if "auto_model" not in kwargs: + kwargs["auto_model"] = AutoModel + + if "add_pooling_layer" not in kwargs: + kwargs["add_pooling_layer"] = False + + model, tokenizer = FastModel.from_pretrained( + model_name=model_name, + max_seq_length=max_seq_length, + dtype=dtype, + load_in_4bit=load_in_4bit, + load_in_8bit=load_in_8bit, + load_in_16bit=load_in_16bit, + full_finetuning=full_finetuning, + token=token, + device_map=device_map, + rope_scaling=rope_scaling, + fix_tokenizer=fix_tokenizer, + trust_remote_code=trust_remote_code, + use_gradient_checkpointing=use_gradient_checkpointing, + resize_model_vocab=resize_model_vocab, + revision=revision, + return_logits=False, + use_exact_model_name=use_exact_model_name, + offload_embedding=offload_embedding, + random_state=random_state, + max_lora_rank=max_lora_rank, + disable_log_stats=disable_log_stats, + qat_scheme=qat_scheme, + load_in_fp8=load_in_fp8, + unsloth_tiled_mlp=unsloth_tiled_mlp, + **kwargs, + ) + + transformer_module = Transformer.__new__(Transformer) + import torch + + torch.nn.Module.__init__(transformer_module) + + transformer_module.auto_model = model + transformer_module.tokenizer = tokenizer + transformer_module.do_lower_case = False + if hasattr(tokenizer, "do_lower_case"): + transformer_module.do_lower_case = tokenizer.do_lower_case + + import inspect + + model_forward_params = list(inspect.signature(model.forward).parameters) + transformer_module.model_forward_params = set(model_forward_params) | { + "input_ids", + "attention_mask", + "token_type_ids", + "inputs_embeds", + } + + if max_seq_length is None: + if ( + hasattr(model, "config") + and hasattr(model.config, "max_position_embeddings") + and hasattr(tokenizer, "model_max_length") + ): + max_seq_length = min( + model.config.max_position_embeddings, tokenizer.model_max_length + ) + elif 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 # default + + 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__ + + hidden_size = model.config.hidden_size + + # detect pooling mode if not specified/default + if pooling_mode == "mean": + try: + from huggingface_hub import hf_hub_download + import json + import os + + if os.path.exists(model_name) and os.path.exists( + os.path.join(model_name, "modules.json") + ): + modules_json_path = os.path.join(model_name, "modules.json") + else: + modules_json_path = hf_hub_download( + model_name, "modules.json", token=token + ) + + with open(modules_json_path, "r") as f: + modules_config = json.load(f) + + pooling_config_path = None + for module in modules_config: + if module.get("type", "") == "sentence_transformers.models.Pooling": + pooling_path = module.get("path", "") + if pooling_path: + # try to find config.json for pooling module + if os.path.exists(model_name) and os.path.exists( + os.path.join(model_name, pooling_path, "config.json") + ): + pooling_config_path = os.path.join( + model_name, pooling_path, "config.json" + ) + else: + pooling_config_path = hf_hub_download( + model_name, + os.path.join(pooling_path, "config.json"), + token=token, + ) + break + + if pooling_config_path: + with open(pooling_config_path, "r") as f: + pooling_config = json.load(f) + if ( + "pooling_mode_cls_token" in pooling_config + and pooling_config["pooling_mode_cls_token"] + ): + print("Pooling mode detected as cls, updating...") + pooling_mode = "cls" + elif ( + "pooling_mode_mean_tokens" in pooling_config + and pooling_config["pooling_mode_mean_tokens"] + ): + print("Pooling mode detected as mean, updating...") + pooling_mode = "mean" + elif ( + "pooling_mode_max_tokens" in pooling_config + and pooling_config["pooling_mode_max_tokens"] + ): + print("Pooling mode detected as max, updating...") + pooling_mode = "max" + elif ( + "pooling_mode_mean_sqrt_len_tokens" in pooling_config + and pooling_config["pooling_mode_mean_sqrt_len_tokens"] + ): + print("Pooling mode detected as mean_sqrt_len, updating...") + pooling_mode = "mean_sqrt_len" + + except Exception as e: + print( + f"Failed to detect pooling mode: {e}, defaulting to mean pooling." + ) + + pooling_module = Pooling( + word_embedding_dimension=hidden_size, + pooling_mode=pooling_mode, + ) + normalize_module = Normalize() + modules = [transformer_module, pooling_module, normalize_module] + st_model = SentenceTransformer(modules=modules) + return st_model + + @staticmethod + def get_peft_model( + model, + r=16, + target_modules=[ + "query", + "key", + "value", + "dense", + ], + lora_alpha=16, + lora_dropout=0.0, + bias="none", + layers_to_transform=None, + layers_pattern=None, + use_gradient_checkpointing="unsloth", + random_state=3407, + max_seq_length=2048, + use_rslora=False, + modules_to_save=None, + init_lora_weights=True, + loftq_config={}, + **kwargs, + ): + from sentence_transformers import SentenceTransformer + + if "task_type" not in kwargs: + kwargs["task_type"] = "FEATURE_EXTRACTION" + print("Setting task_type to FEATURE_EXTRACTION") + + if isinstance(model, SentenceTransformer): + # extract inner model from the transformer module + transformer_module = model[0] + inner_model = transformer_module.auto_model + + peft_model = FastModel.get_peft_model( + model=inner_model, + r=r, + target_modules=target_modules, + lora_alpha=lora_alpha, + lora_dropout=lora_dropout, + bias=bias, + layers_to_transform=layers_to_transform, + layers_pattern=layers_pattern, + use_gradient_checkpointing=use_gradient_checkpointing, + random_state=random_state, + max_seq_length=max_seq_length, + use_rslora=use_rslora, + modules_to_save=modules_to_save, + init_lora_weights=init_lora_weights, + loftq_config=loftq_config, + **kwargs, + ) + + # re-assign the peft model back to the transformer module + transformer_module.auto_model = peft_model + return model + else: + return FastModel.get_peft_model( + model=model, + r=r, + target_modules=target_modules, + lora_alpha=lora_alpha, + lora_dropout=lora_dropout, + bias=bias, + layers_to_transform=layers_to_transform, + layers_pattern=layers_pattern, + use_gradient_checkpointing=use_gradient_checkpointing, + random_state=random_state, + max_seq_length=max_seq_length, + use_rslora=use_rslora, + modules_to_save=modules_to_save, + init_lora_weights=init_lora_weights, + loftq_config=loftq_config, + **kwargs, + )