unslothai--unsloth/unsloth/models/sentence_transformer.py

# 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.

import logging

from .loader import FastModel, DISABLE_SDPA_MODEL_NAMES
from ._utils import (
    SUPPORTS_BFLOAT16,
    resolve_model_class,
    resolve_encoder_attention_implementation,
    maybe_prefetch_hf_snapshot,
)
import inspect
import json
import os
import threading
import types
from huggingface_hub import hf_hub_download
from typing import Optional
import torch
from transformers.modeling_outputs import BaseModelOutput
from collections import OrderedDict
from transformers.models.distilbert import modeling_distilbert
from transformers.modeling_attn_mask_utils import _prepare_4d_attention_mask_for_sdpa
import transformers
from packaging.version import Version
import re
from transformers import AutoModel, AutoConfig
import tempfile
from huggingface_hub import HfApi, get_token
from ..save import unsloth_save_pretrained_torchao, unsloth_save_pretrained_gguf
import contextlib
import shutil


_CREATE_TRANSFORMER_MODULE_LOCK = threading.RLock()


def _save_pretrained_torchao(
    self,
    save_directory,
    tokenizer = None,
    torchao_config = None,
    push_to_hub = False,
    token = None,
):
    self.save_pretrained(save_directory)

    inner_model = self[0].auto_model
    if hasattr(inner_model, "_orig_mod"):
        inner_model = inner_model._orig_mod

    # merge LoRA first
    if hasattr(inner_model, "merge_and_unload"):
        inner_model = inner_model.merge_and_unload()

    transformer_path = "0_Transformer"
    modules_path = os.path.join(save_directory, "modules.json")
    if os.path.exists(modules_path):
        try:
            with open(modules_path, "r", encoding = "utf-8") as f:
                modules = json.load(f)
            for m in modules:
                if m.get("type", "").endswith("Transformer"):
                    transformer_path = m.get("path", "")
                    break
        except:
            pass

    transformer_dir = os.path.join(save_directory, transformer_path)
    transformer_dir = os.path.abspath(transformer_dir)

    if tokenizer is None:
        tokenizer = self.tokenizer

    @contextlib.contextmanager
    def patch_unsloth_save():
        original_causal = transformers.AutoModelForCausalLM
        original_rmtree = shutil.rmtree
        # unsloth_save_pretrained_torchao expects AutoModelForCausalLM
        transformers.AutoModelForCausalLM = transformers.AutoModel
        # prevent unsloth from deleting the unquantized model directory
        shutil.rmtree = lambda *args, **kwargs: None
        try:
            yield
        finally:
            transformers.AutoModelForCausalLM = original_causal
            shutil.rmtree = original_rmtree

    with patch_unsloth_save():
        unsloth_save_pretrained_torchao(
            inner_model,
            transformer_dir,
            tokenizer = tokenizer,
            torchao_config = torchao_config,
            push_to_hub = push_to_hub,
            token = token,
        )

    # avoid `0_Transformer-torchao`, it was either this or fix modules.json
    torchao_dir = transformer_dir + "-torchao"
    if os.path.exists(torchao_dir):
        if not os.path.exists(transformer_dir):
            os.makedirs(transformer_dir, exist_ok = True)

        for item in os.listdir(torchao_dir):
            s = os.path.join(torchao_dir, item)
            d = os.path.join(transformer_dir, item)
            if os.path.isdir(s):
                shutil.copytree(s, d, dirs_exist_ok = True)
            else:
                shutil.copy2(s, d)

        shutil.rmtree(torchao_dir)

        # remove conflicting safetensors if we brought in bin
        if os.path.exists(os.path.join(transformer_dir, "pytorch_model.bin")):
            safetensors_path = os.path.join(transformer_dir, "model.safetensors")
            if os.path.exists(safetensors_path):
                try:
                    os.remove(safetensors_path)
                except:
                    pass

    try:
        FastSentenceTransformer._add_unsloth_branding(save_directory)
    except:
        pass


# Thanks Etherl:
def _save_pretrained_gguf(
    self,
    save_directory,
    tokenizer = None,
    quantization_method = "fast_quantized",
    first_conversion = None,
    push_to_hub = False,
    token = None,
    max_shard_size = "5GB",
    temporary_location = "_unsloth_temporary_saved_buffers",
    maximum_memory_usage = 0.85,
    **kwargs,
):
    """
    Saves the SentenceTransformer model to GGUF format by saving the inner transformer model,
    converting it, and placing the resulting GGUF files in the save directory.
    """
    # 1. Save standard SentenceTransformer structure (configs, modules.json, etc.)
    self.save_pretrained(save_directory)

    # 2. Extract inner transformer model (PEFT merge handled by the gguf saver)
    inner_model = self[0].auto_model
    if hasattr(inner_model, "_orig_mod"):
        inner_model = inner_model._orig_mod

    # 3. Identify where the transformer weights are stored
    transformer_path = "0_Transformer"
    modules_path = os.path.join(save_directory, "modules.json")
    if os.path.exists(modules_path):
        try:
            with open(modules_path, "r", encoding = "utf-8") as f:
                modules = json.load(f)
            for m in modules:
                if m.get("type", "").endswith("Transformer"):
                    transformer_path = m.get("path", "")
                    break
        except:
            pass

    # Unsloth saves + converts here; absolute for later commonpath comparison
    transformer_dir = os.path.join(save_directory, transformer_path)
    transformer_dir = os.path.abspath(transformer_dir)

    if tokenizer is None:
        tokenizer = self.tokenizer

    # 4. Patch environment so Unsloth treats this embedding model correctly
    @contextlib.contextmanager
    def patch_unsloth_gguf_save():
        # Prevent deletion of the directory self.save_pretrained just created
        original_rmtree = shutil.rmtree
        try:
            yield
        finally:
            shutil.rmtree = original_rmtree

    # 5. Call Unsloth's GGUF saver on the inner model targeting the transformer subdirectory
    with patch_unsloth_gguf_save():
        result = unsloth_save_pretrained_gguf(
            inner_model,
            save_directory = transformer_dir,
            tokenizer = tokenizer,
            quantization_method = quantization_method,
            first_conversion = first_conversion,
            push_to_hub = False,  # Force local first to move files
            token = token,
            max_shard_size = max_shard_size,
            temporary_location = temporary_location,
            maximum_memory_usage = maximum_memory_usage,
        )

    # 6. Move GGUF files from the subdirectory (0_Transformer) to the root save_directory
    gguf_files = result.get("gguf_files", [])

    new_gguf_locations = []

    for gguf_file in gguf_files:
        if os.path.exists(gguf_file):
            filename = os.path.basename(gguf_file)
            dest_path = os.path.join(save_directory, filename)

            abs_gguf_file = os.path.abspath(gguf_file)

            try:
                is_subpath = os.path.commonpath([abs_gguf_file, transformer_dir]) == transformer_dir
            except ValueError:
                # Windows different-drive commonpath raises
                is_subpath = False

            if is_subpath:
                # Move the GGUF out of transformer_dir to root
                shutil.move(gguf_file, dest_path)
                new_gguf_locations.append(dest_path)
            else:
                # Elsewhere: move to root if not already there
                if os.path.abspath(dest_path) != abs_gguf_file:
                    shutil.move(gguf_file, dest_path)
                new_gguf_locations.append(dest_path)

    result["gguf_files"] = new_gguf_locations

    # 7. Add branding
    try:
        FastSentenceTransformer._add_unsloth_branding(save_directory)

        # Add GGUF details to README
        readme_path = os.path.join(save_directory, "README.md")
        if os.path.exists(readme_path):
            with open(readme_path, "a", encoding = "utf-8") as f:
                f.write("\n## GGUF Quantization\n")
                f.write(
                    f"This model contains GGUF quantized versions in: {', '.join([os.path.basename(f) for f in new_gguf_locations])}\n"
                )
    except:
        pass

    # 8. Handle Push to Hub if requested
    if push_to_hub:
        if token is None:
            token = get_token()

        api = HfApi(token = token)
        repo_id = save_directory  # Assuming save_directory is the repo name if pushing

        print(f"Unsloth: Uploading to {repo_id}...")
        try:
            api.create_repo(repo_id = repo_id, exist_ok = True, private = kwargs.get("private", False))
            api.upload_folder(
                folder_path = save_directory,
                repo_id = repo_id,
                commit_message = "Upload GGUF and SentenceTransformer model",
            )
            print(f"Unsloth: Uploaded to https://huggingface.co/{repo_id}")
        except Exception as e:
            print(f"Unsloth: Upload failed: {e}")

    return result


def _push_to_hub_gguf(
    self,
    repo_id,
    tokenizer = None,
    quantization_method = "fast_quantized",
    first_conversion = None,
    token = None,
    private = None,
    commit_message = "Upload GGUF SentenceTransformer model trained with Unsloth",
    commit_description = "Upload GGUF model trained with Unsloth 2x faster",
    max_shard_size = "5GB",
    temporary_location = "_unsloth_temporary_saved_buffers",
    maximum_memory_usage = 0.85,
    create_pr = False,
    revision = None,
    tags = None,
    **kwargs,
):
    """
    Converts the SentenceTransformer model to GGUF format and pushes to the Hugging Face Hub.

    This method:
    1. Saves the model locally to a temporary directory in GGUF format.
    2. Uploads the GGUF files, config, Ollama Modelfile, and README to the Hub.
    3. Cleans up the temporary directory.

    Args:
        repo_id (str): The Hugging Face Hub repo ID (e.g., "username/model-name").
        tokenizer: The tokenizer to save. Defaults to `self.tokenizer`.
        quantization_method (str or list): GGUF quantization method(s). Can be a string or list of strings.
            Choose from the following options:
            * "not_quantized"  : Recommended. Fast conversion. Slow inference, big files.
            * "fast_quantized" : Recommended. Fast conversion. OK inference, OK file size.
            * "quantized"      : Recommended. Slow conversion. Fast inference, small files.
            * "f32"     : Not recommended. Retains 100% accuracy, but super slow and memory hungry.
            * "f16"     : Fastest conversion + retains 100% accuracy. Slow and memory hungry.
            * "q8_0"    : Fast conversion. High resource use, but generally acceptable.
            * "q4_k_m"  : Recommended. Uses Q6_K for half of the attention.wv and feed_forward.w2 tensors, else Q4_K
            * "q5_k_m"  : Recommended. Uses Q6_K for half of the attention.wv and feed_forward.w2 tensors, else Q5_K
            * "q2_k"    : Uses Q4_K for the attention.vw and feed_forward.w2 tensors, Q2_K for the other tensors.
            * "q3_k_l"  : Uses Q5_K for the attention.wv, attention.wo, and feed_forward.w2 tensors, else Q3_K
            * "q3_k_m"  : Uses Q4_K for the attention.wv, attention.wo, and feed_forward.w2 tensors, else Q3_K
            * "q3_k_s"  : Uses Q3_K for all tensors
            * "q4_0"    : Original quant method, 4-bit.
            * "q4_1"    : Higher accuracy than q4_0 but not as high as q5_0. However has quicker inference than q5 models.
            * "q4_k_s"  : Uses Q4_K for all tensors
            * "q5_0"    : Higher accuracy, higher resource usage and slower inference.
            * "q5_1"    : Even higher accuracy, resource usage and slower inference.
            * "q5_k_s"  : Uses Q5_K for all tensors
            * "q6_k"    : Uses Q8_K for all tensors
        first_conversion (str, optional): The initial conversion format before quantization.
        token (str, optional): Hugging Face token. Uses cached token if not provided.
        private (bool, optional): Whether the repo should be private.
        commit_message (str): Commit message for the upload.
        commit_description (str): Commit description for the upload.
        max_shard_size (str): Maximum shard size for saving.
        temporary_location (str): Temp directory for intermediate files.
        maximum_memory_usage (float): Max fraction of memory to use.
        create_pr (bool): Whether to create a pull request instead of pushing directly.
        revision (str, optional): Branch/revision to push to.
        tags (list, optional): Additional tags for the repo.

    Returns:
        str: The full repo ID on Hugging Face Hub.
    """
    if token is None:
        token = get_token()
    if token is None:
        raise ValueError(
            "No HF token provided. Please provide a token or login with `huggingface-cli login`"
        )

    api = HfApi(token = token)

    if "/" not in repo_id:
        username = api.whoami()["name"]
        full_repo_id = f"{username}/{repo_id}"
    else:
        full_repo_id = repo_id

    model_name = full_repo_id.split("/")[-1]

    try:
        api.create_repo(
            repo_id = full_repo_id,
            private = private,
            exist_ok = True,
            repo_type = "model",
        )
    except Exception as e:
        print(f"Unsloth Warning: Could not create repo: {e}")

    # Convert locally in a temp dir, then upload
    with tempfile.TemporaryDirectory(prefix = "unsloth_st_gguf_") as temp_dir:
        print(f"Unsloth: Converting SentenceTransformer to GGUF format...")

        result = _save_pretrained_gguf(
            self,
            save_directory = temp_dir,
            tokenizer = tokenizer,
            quantization_method = quantization_method,
            first_conversion = first_conversion,
            push_to_hub = False,  # We handle upload ourselves
            token = token,
            max_shard_size = max_shard_size,
            temporary_location = temporary_location,
            maximum_memory_usage = maximum_memory_usage,
        )

        gguf_files = result.get("gguf_files", [])
        modelfile_location = result.get("modelfile_location", None)
        is_vlm = result.get("is_vlm", False)
        fix_bos_token = result.get("fix_bos_token", False)

        print(f"Unsloth: Uploading GGUF to https://huggingface.co/{full_repo_id}...")

        # Upload GGUF files
        for file_location in gguf_files:
            if os.path.exists(file_location):
                filename = os.path.basename(file_location)
                print(f"  Uploading {filename}...")
                api.upload_file(
                    path_or_fileobj = file_location,
                    path_in_repo = filename,
                    repo_id = full_repo_id,
                    repo_type = "model",
                    commit_message = commit_message,
                    commit_description = commit_description,
                    create_pr = create_pr,
                    revision = revision,
                )

        # Upload Modelfile if exists
        if modelfile_location and os.path.exists(modelfile_location):
            print("  Uploading Ollama Modelfile...")
            api.upload_file(
                path_or_fileobj = modelfile_location,
                path_in_repo = "Modelfile",
                repo_id = full_repo_id,
                repo_type = "model",
                commit_message = f"{commit_message} - Ollama Modelfile",
                create_pr = create_pr,
                revision = revision,
            )

        # Upload config.json if exists
        config_path = os.path.join(temp_dir, "config.json")
        if os.path.exists(config_path):
            print("  Uploading config.json...")
            api.upload_file(
                path_or_fileobj = config_path,
                path_in_repo = "config.json",
                repo_id = full_repo_id,
                repo_type = "model",
                commit_message = f"{commit_message} - config",
                create_pr = create_pr,
                revision = revision,
            )

        # Create and upload README
        gguf_basenames = [os.path.basename(f) for f in gguf_files if os.path.exists(f)]
        readme_content = f"""---
tags:
- gguf
- llama.cpp
- unsloth
- sentence-transformers
{"- vision-language-model" if is_vlm else ""}
---

# {model_name} - GGUF

This sentence-transformers model was finetuned and converted to GGUF format using [Unsloth](https://github.com/unslothai/unsloth).

## Available Model files:
"""
        for fname in gguf_basenames:
            readme_content += f"- `{fname}`\n"

        if modelfile_location and os.path.exists(modelfile_location):
            readme_content += "\n## Ollama\n"
            readme_content += "An Ollama Modelfile is included for easy deployment.\n"

        if fix_bos_token:
            readme_content += "\n## Note\n"
            readme_content += (
                "The model's BOS token behavior was adjusted for GGUF compatibility.\n"
            )

        readme_content += (
            "\nThis was trained 2x faster with [Unsloth](https://github.com/unslothai/unsloth)\n"
            '[<img src="https://raw.githubusercontent.com/unslothai/unsloth/main/images/unsloth%20made%20with%20love.png" width="200"/>](https://github.com/unslothai/unsloth)\n'
        )

        readme_path = os.path.join(temp_dir, "README.md")
        with open(readme_path, "w", encoding = "utf-8") as f:
            f.write(readme_content)

        api.upload_file(
            path_or_fileobj = readme_path,
            path_in_repo = "README.md",
            repo_id = full_repo_id,
            repo_type = "model",
            commit_message = "Add README",
            create_pr = create_pr,
            revision = revision,
        )

    # Add tags
    all_tags = ["gguf", "llama-cpp", "unsloth", "sentence-transformers"]
    if is_vlm:
        all_tags.append("vision-language-model")
    if tags is not None:
        if isinstance(tags, (list, tuple)):
            all_tags.extend(tags)
        else:
            all_tags.append(tags)
    try:
        api.add_tags(repo_id = full_repo_id, tags = all_tags, repo_type = "model")
    except:
        pass

    print(f"Unsloth: Successfully uploaded GGUF to https://huggingface.co/{full_repo_id}")
    return full_repo_id


class FastSentenceTransformer(FastModel):
    @staticmethod
    def _save_base_config_for_processor_resume(config, output_path):
        """sentence-transformers >= 5.4 reloads Transformer modules via
        AutoProcessor, which falls back to AutoConfig for tokenizer-only
        roots -- so PEFT adapter checkpoints still need base config.json
        next to adapter_config.json."""
        if config is None or not getattr(config, "model_type", None):
            return
        if hasattr(config, "save_pretrained"):
            config.save_pretrained(output_path)
        elif hasattr(config, "to_json_file"):
            config_path = os.path.join(output_path, "config.json")
            config.to_json_file(config_path)

    @staticmethod
    def _patch_transformer_module_save_config(transformer_module, base_config = None):
        transformer_module._unsloth_st_managed = True
        if base_config is not None and getattr(base_config, "model_type", None):
            transformer_module._unsloth_base_config = base_config

        if getattr(transformer_module, "_unsloth_save_config_patched", False):
            return transformer_module

        original_save = transformer_module.save

        def _save_with_base_config(self, output_path, *args, **kwargs):
            original_save(output_path, *args, **kwargs)
            FastSentenceTransformer._save_base_config_for_processor_resume(
                getattr(self, "_unsloth_base_config", None), output_path
            )

        transformer_module.save = types.MethodType(_save_with_base_config, transformer_module)
        transformer_module._unsloth_save_config_patched = True
        return transformer_module

    @staticmethod
    def _read_pooling_mode(
        model_name,
        token,
        cache_dir = None,
        revision = None,
    ):
        """Read the pooling mode from modules.json, else return "mean"."""
        try:
            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,
                    cache_dir = cache_dir,
                    revision = revision,
                )

            with open(modules_json_path, "r", encoding = "utf-8") 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:
                        # find config.json for the 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,
                                cache_dir = cache_dir,
                                revision = revision,
                            )
                        break

            if pooling_config_path:
                with open(pooling_config_path, "r", encoding = "utf-8") as f:
                    pooling_config = json.load(f)
                    # from here:
                    # https://github.com/huggingface/sentence-transformers/blob/main/sentence_transformers/models/Pooling.py#L43
                    pooling_map = {
                        "pooling_mode_cls_token": "cls",
                        "pooling_mode_mean_tokens": "mean",
                        "pooling_mode_max_tokens": "max",
                        "pooling_mode_mean_sqrt_len_tokens": "mean_sqrt_len",
                        "pooling_mode_weightedmean_tokens": "weightedmean",
                        "pooling_mode_lasttoken": "lasttoken",
                    }
                    for config_key, mode in pooling_map.items():
                        if pooling_config.get(config_key):
                            if mode != "mean":
                                print(f"Pooling mode detected as {mode}, updating...")
                            return mode

        except Exception as e:
            print(
                f"Failed to detect pooling mode, not a sentence-transformers model. Using default pooling mode 'mean', this may or may not work."
            )
            return "mean"

    # should prolly be done upstream instead of this hackfest here
    @staticmethod
    def _patch_mpnet_v4():
        """Patch MPNetModel for gradient checkpointing (transformers 4)."""
        from transformers.models.mpnet import modeling_mpnet

        modeling_mpnet.MPNetModel.supports_gradient_checkpointing = True

        def _set_gradient_checkpointing(
            self,
            module = None,
            value = True,
        ):
            if module is None:
                module = self.encoder
            if isinstance(module, modeling_mpnet.MPNetEncoder):
                module.gradient_checkpointing = value

        modeling_mpnet.MPNetModel._set_gradient_checkpointing = _set_gradient_checkpointing

        # patch MPNetEncoder.forward for checkpointing; based on:
        # https://github.com/huggingface/transformers/blob/v4.57.3/src/transformers/models/mpnet/modeling_mpnet.py#L321
        def forward(
            self,
            hidden_states: torch.Tensor,
            attention_mask: Optional[torch.Tensor] = None,
            head_mask: Optional[torch.Tensor] = None,
            output_attentions: bool = False,
            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

            for i, layer_module in enumerate(self.layer):
                if output_hidden_states:
                    all_hidden_states = all_hidden_states + (hidden_states,)

                if getattr(self, "gradient_checkpointing", False) and self.training:

                    def create_custom_forward(module):
                        def custom_forward(*inputs):
                            return module(*inputs, output_attentions = output_attentions)

                        return custom_forward

                    layer_outputs = torch.utils.checkpoint.checkpoint(
                        create_custom_forward(layer_module),
                        hidden_states,
                        attention_mask,
                        head_mask[i] if head_mask is not None else None,
                        position_bias,
                        use_reentrant = True,  # fix for torch 2.9
                    )
                else:
                    layer_outputs = layer_module(
                        hidden_states,
                        attention_mask,
                        head_mask[i] if head_mask is not None else None,
                        position_bias,
                        output_attentions = output_attentions,
                        **kwargs,
                    )

                hidden_states = layer_outputs[0]

                if output_attentions:
                    all_attentions = all_attentions + (layer_outputs[1],)

            if output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states,)

            if not return_dict:
                return tuple(
                    v for v in [hidden_states, all_hidden_states, all_attentions] if v is not None
                )
            return BaseModelOutput(
                last_hidden_state = hidden_states,
                hidden_states = all_hidden_states,
                attentions = all_attentions,
            )

        modeling_mpnet.MPNetEncoder.forward = forward

    @staticmethod
    def _patch_mpnet_v5():
        """Patch MPNetModel for gradient checkpointing (transformers 5)."""
        from transformers.models.mpnet import modeling_mpnet

        modeling_mpnet.MPNetModel.supports_gradient_checkpointing = True

        def _set_gradient_checkpointing(
            self,
            module = None,
            value = True,
        ):
            if module is None:
                module = self.encoder
            if isinstance(module, modeling_mpnet.MPNetEncoder):
                module.gradient_checkpointing = value

        modeling_mpnet.MPNetModel._set_gradient_checkpointing = _set_gradient_checkpointing

        # patch MPNetEncoder.forward for checkpointing; based on:
        # https://github.com/huggingface/transformers/blob/v5.0.0rc1/src/transformers/models/mpnet/modeling_mpnet.py#L284
        def forward(
            self,
            hidden_states: torch.Tensor,
            attention_mask: Optional[torch.Tensor] = None,
            output_attentions: bool = False,
            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

            for i, layer_module in enumerate(self.layer):
                if output_hidden_states:
                    all_hidden_states = all_hidden_states + (hidden_states,)

                if getattr(self, "gradient_checkpointing", False) and self.training:

                    def create_custom_forward(module):
                        def custom_forward(*inputs):
                            return module(*inputs, output_attentions = output_attentions)

                        return custom_forward

                    layer_outputs = torch.utils.checkpoint.checkpoint(
                        create_custom_forward(layer_module),
                        hidden_states,
                        attention_mask,
                        position_bias,
                        use_reentrant = True,  # required for torch >= 2.9
                    )
                else:
                    layer_outputs = layer_module(
                        hidden_states,
                        attention_mask,
                        position_bias,
                        output_attentions,
                        **kwargs,
                    )

                hidden_states = layer_outputs[0]

                if output_attentions:
                    all_attentions = all_attentions + (layer_outputs[1],)

            if output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states,)

            if not return_dict:
                return tuple(
                    v for v in [hidden_states, all_hidden_states, all_attentions] if v is not None
                )
            return BaseModelOutput(
                last_hidden_state = hidden_states,
                hidden_states = all_hidden_states,
                attentions = all_attentions,
            )

        modeling_mpnet.MPNetEncoder.forward = forward

    @staticmethod
    def _patch_distilbert_v4():
        # change kwargs to positional args to be compatible with peft_utils
        """Patch DistilBertModel.forward to use positional args (transformers 4)."""

        # based on:
        # https://github.com/huggingface/transformers/blob/v4.57.3/src/transformers/models/distilbert/modeling_distilbert.py#L666
        def forward(
            self,
            input_ids: Optional[torch.Tensor] = None,
            attention_mask: Optional[torch.Tensor] = None,
            head_mask: Optional[torch.Tensor] = None,
            inputs_embeds: Optional[torch.Tensor] = None,
            output_attentions: Optional[bool] = None,
            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
            )
            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"
                )
            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")

            device = input_ids.device if input_ids is not None else inputs_embeds.device

            head_mask_is_none = head_mask is None
            head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)

            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
                )
            else:
                if attention_mask is None:
                    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
                ):
                    attention_mask = _prepare_4d_attention_mask_for_sdpa(
                        attention_mask, embeddings.dtype, tgt_len = input_shape[1]
                    )
            # patch here, change kwargs to positional args:
            return self.transformer(
                embeddings,
                attention_mask,
                head_mask,
                output_attentions,
                output_hidden_states,
                return_dict,
            )

        modeling_distilbert.DistilBertModel.forward = forward

    @staticmethod
    def _has_add_pooling_layer(config, auto_model_class = None):
        """Check if the model class accepts the `add_pooling_layer` argument."""
        try:
            if auto_model_class is None:
                auto_model_class = AutoModel
            model_class = resolve_model_class(auto_model_class, config)

            if model_class:
                sig = inspect.signature(model_class.__init__)
                return "add_pooling_layer" in sig.parameters
        except:
            pass

        return False

    @staticmethod
    def _patch_distilbert_v5():
        """Patch DistilBertModel.forward to use positional args (transformers 5)."""
        # based on:
        # https://github.com/huggingface/transformers/blob/v5.0.0rc1/src/transformers/models/distilbert/modeling_distilbert.py#L386
        from transformers.masking_utils import create_bidirectional_mask

        def forward(
            self,
            input_ids: Optional[torch.Tensor] = None,
            attention_mask: Optional[torch.Tensor] = None,
            inputs_embeds: Optional[torch.Tensor] = None,
            position_ids: Optional[torch.Tensor] = None,
            **kwargs,
        ):
            if (input_ids is None) ^ (inputs_embeds is not None):
                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,
            )

            # patch here: unsloth gradient checkpointing hook needs positional arguments
            return self.transformer(
                embeddings,
                attention_mask,
                **kwargs,
            )

        modeling_distilbert.DistilBertModel.forward = forward

    @staticmethod
    def _add_unsloth_tags(
        repo_id,
        token,
        tags = None,
    ):
        """Add Unsloth + sentence-transformers tags to the HF Hub repo."""
        from huggingface_hub import HfApi

        api = HfApi(token = token)
        if tags is None:
            tags = []
        tags.extend(["unsloth", "sentence-transformers"])
        try:
            api.add_tags(
                repo_id = repo_id,
                tags = tags,
                repo_type = "model",
            )
        except:
            pass

    @staticmethod
    def _add_unsloth_branding(save_directory):
        """Add Unsloth branding to the sentence-transformers-generated README.md."""
        readme_path = os.path.join(save_directory, "README.md")
        if not os.path.exists(readme_path):
            return

        with open(readme_path, "r", encoding = "utf-8") as f:
            content = f.read()

        # add unsloth tag to frontmatter
        if "---\ntags:\n" in content:
            content = content.replace("---\ntags:\n", "---\ntags:\n- unsloth\n")
        else:
            # tags exist but not at the start: append via regex
            pattern = r"(^tags:\s*\n)"
            if re.search(pattern, content, re.MULTILINE):
                content = re.sub(pattern, r"\1- unsloth\n", content, count = 1, flags = re.MULTILINE)

        branding = (
            "\n\nThis model was finetuned with [Unsloth](https://github.com/unslothai/unsloth).\n\n"
            '[<img src="https://raw.githubusercontent.com/unslothai/unsloth/main/images/unsloth%20made%20with%20love.png" width="200"/>](https://github.com/unslothai/unsloth)\n'
        )

        if "# SentenceTransformer" in content:
            parts = content.split("# SentenceTransformer", 1)
            content = parts[0] + "# SentenceTransformer" + branding + parts[1]
        else:
            content += branding

        with open(readme_path, "w", encoding = "utf-8") as f:
            f.write(content)

    @staticmethod
    def _module_path(
        model_name,
        token = None,
        cache_dir = None,
        revision = None,
    ):
        """Return the path to the modules.json file, or None."""
        try:
            if os.path.exists(model_name) and os.path.isdir(model_name):
                path = os.path.join(model_name, "modules.json")
                return path if os.path.exists(path) else None
            else:
                try:
                    return hf_hub_download(
                        model_name,
                        "modules.json",
                        token = token,
                        cache_dir = cache_dir,
                        revision = revision,
                    )
                except:
                    return None
        except:
            return None

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

        # Prevents loading the model a second time and redirects AutoProcessor/
        # AutoTokenizer so Transformer.__init__ picks up our pre-fixed tokenizer.
        # On sentence-transformers >=5.4 `tokenizer` is a read-only @property
        # backed by `self.processor`, so a post-init assignment raises; the
        # constructor redirect sets self.processor correctly instead.
        from transformers import AutoProcessor, AutoTokenizer

        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) as exception:
                logging.debug(
                    "Unsloth: Could not normalize SentenceTransformer model path: %s",
                    exception,
                )
                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) as exception:
                logging.debug(
                    "Unsloth: Could not compare SentenceTransformer model paths: %s",
                    exception,
                )
            return False

        with _CREATE_TRANSFORMER_MODULE_LOCK:
            original_model_from_pretrained = AutoModel.from_pretrained
            original_processor_from_pretrained = AutoProcessor.from_pretrained
            original_tokenizer_from_pretrained = AutoTokenizer.from_pretrained

            def return_existing_model(*args, **kwargs):
                if is_requested_model_name(args, kwargs):
                    return model
                return original_model_from_pretrained(*args, **kwargs)

            def return_existing_tokenizer(*args, **kwargs):
                if is_requested_model_name(args, kwargs):
                    return tokenizer
                return original_tokenizer_from_pretrained(*args, **kwargs)

            def return_existing_processor(*args, **kwargs):
                if is_requested_model_name(args, kwargs):
                    return tokenizer
                return original_processor_from_pretrained(*args, **kwargs)

            try:
                # Temporarily redirect Auto* loading to return our pre-loaded objects
                AutoModel.from_pretrained = return_existing_model
                AutoProcessor.from_pretrained = return_existing_processor
                AutoTokenizer.from_pretrained = return_existing_tokenizer

                transformer_init_params = inspect.signature(Transformer.__init__).parameters
                trust_remote_code_kwargs = {"trust_remote_code": trust_remote_code}
                do_lower_case = getattr(tokenizer, "do_lower_case", False)
                transformer_kwargs = {"max_seq_length": max_seq_length}
                if "do_lower_case" in transformer_init_params:
                    transformer_kwargs["do_lower_case"] = do_lower_case
                if "model_kwargs" in transformer_init_params:
                    transformer_kwargs["model_kwargs"] = trust_remote_code_kwargs.copy()
                    transformer_kwargs["config_kwargs"] = trust_remote_code_kwargs.copy()
                else:
                    transformer_kwargs["model_args"] = trust_remote_code_kwargs.copy()
                    transformer_kwargs["config_args"] = trust_remote_code_kwargs.copy()
                if "processor_kwargs" in transformer_init_params:
                    transformer_kwargs["processor_kwargs"] = trust_remote_code_kwargs.copy()
                elif "tokenizer_args" in transformer_init_params:
                    transformer_kwargs["tokenizer_args"] = trust_remote_code_kwargs.copy()

                # 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
                AutoProcessor.from_pretrained = original_processor_from_pretrained
                AutoTokenizer.from_pretrained = original_tokenizer_from_pretrained

        # On sentence-transformers >=5.4 `tokenizer` is a read-only property backed
        # by `self.processor` (already wired via the redirect above). On older
        # versions it's a regular attribute and the explicit assignment is required.
        if not isinstance(getattr(type(transformer_module), "tokenizer", None), property):
            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
        preinit_model_forward_params = getattr(transformer_module, "model_forward_params", set())
        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",
            "return_dict",
        }
        transformer_module.model_forward_params |= preinit_model_forward_params

        # determine max_seq_length if not provided
        if max_seq_length is None:
            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
        config_keys = list(getattr(transformer_module, "config_keys", []) or [])
        for config_key in ("max_seq_length", "do_lower_case"):
            if config_key not in config_keys:
                config_keys.append(config_key)
        transformer_module.config_keys = config_keys
        transformer_module.save_in_root = True
        FastSentenceTransformer._patch_transformer_module_save_config(
            transformer_module, getattr(model, "config", None)
        )

        if hasattr(model, "config"):
            model.config.tokenizer_class = tokenizer.__class__.__name__

        return transformer_module

    @staticmethod
    def _is_transformer_module_ref(class_ref):
        if class_ref in {
            "sentence_transformers.models.Transformer",
            "sentence_transformers.models.transformer.Transformer",
            "sentence_transformers.base.modules.transformer.Transformer",
        }:
            return True

        try:
            from sentence_transformers.models import Transformer
            from sentence_transformers.util import import_from_string

            module_class = import_from_string(class_ref)
            return module_class is Transformer
        except (ImportError, AttributeError, TypeError, ValueError) as exception:
            logging.debug(
                "Unsloth: Could not resolve SentenceTransformer module ref %r: %s",
                class_ref,
                exception,
            )
            return False

    @staticmethod
    def _load_modules(
        model_name,
        token,
        model,
        tokenizer,
        max_seq_length,
        pooling_mode,
        trust_remote_code = False,
        cache_dir = None,
        revision = None,
    ) -> tuple[OrderedDict, bool]:
        """Load modules from modules.json, else fall back 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, cache_dir = cache_dir, revision = revision
        )

        if modules_json_path:
            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))))

                if FastSentenceTransformer._is_transformer_module_ref(class_ref):
                    transformer_module = FastSentenceTransformer._create_transformer_module(
                        model_name,
                        model,
                        tokenizer,
                        max_seq_length,
                        trust_remote_code,
                        token,
                        cache_dir,
                        revision,
                        module_subfolder = module_config.get("path") or "",
                    )
                    modules[name] = transformer_module
                else:
                    # load other modules (Pooling, Normalize, etc.)
                    module_path = module_config["path"]
                    if os.path.isdir(model_name):
                        load_path = os.path.join(model_name, module_path)
                    else:
                        try:
                            load_path = load_dir_path(
                                model_name,
                                module_path,
                                token = token,
                                cache_folder = cache_dir,
                                revision = revision,
                            )
                        except Exception as e:
                            print(f"Unsloth Warning: Could not download module {module_path}: {e}")
                            continue

                    module_class = import_from_string(class_ref)
                    try:
                        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}")

            return modules, False

        # fallback if no modules.json (non sentence-transformers models)
        print(
            "Unsloth: No modules.json found, falling back to [Transformer, Pooling, Normalize]. This may or may not work."
        )

        transformer_module = FastSentenceTransformer._create_transformer_module(
            model_name,
            model,
            tokenizer,
            max_seq_length,
            trust_remote_code,
            token,
            cache_dir,
            revision,
        )
        modules["0"] = transformer_module

        hidden_size = getattr(model.config, "hidden_size", 768)

        if pooling_mode == "mean":
            pooling_mode = FastSentenceTransformer._read_pooling_mode(
                model_name, token, cache_dir = cache_dir, revision = revision
            )

        modules["1"] = Pooling(word_embedding_dimension = hidden_size, pooling_mode = pooling_mode)
        modules["2"] = Normalize()

        return modules, True

    # Encoder model types that benefit from native torch.compile instead of Unsloth patching
    ENCODER_MODEL_TYPES = {
        "mpnet",
        "bert",
        "distilbert",
        "modernbert",
        "roberta",
        "xlm-roberta",
        "albert",
        "electra",
    }

    @staticmethod
    def _estimate_compile_threshold(
        model,
        batch_size = None,
        grad_accum = None,
        max_seq_length = None,
    ):
        """Estimate the minimum training steps for torch.compile to pay off
        (with a 1.2x safety margin), from empirical benchmarks.

        Optional batch_size / grad_accum / max_seq_length give a coarse,
        conservative pre-run adjustment with no runtime measurements.
        """
        if hasattr(model, "__getitem__"):
            try:
                inner = model[0].auto_model
                params = sum(p.numel() for p in inner.parameters())
            except:
                params = 100_000_000  # Default to 100M if can't determine
        else:
            params = sum(p.numel() for p in model.parameters())

        model_type = None
        try:
            if "inner" in locals():
                model_type = getattr(getattr(inner, "config", None), "model_type", None)
        except Exception:
            model_type = None
        if isinstance(model_type, str):
            model_type = model_type.lower()

        params_m = params / 1e6

        # Empirical formula based on benchmarks with batch_size=2, grad_accum=4
        # Small models: high fixed overhead, lower speedup
        # Large models: warmup scales but speedup is significant
        if params_m < 50:
            estimated_warmup = 35 + params_m * 0.3
            base_speedup = 1.35
        elif params_m < 200:
            estimated_warmup = 12 + params_m * 0.03
            base_speedup = 1.75
        else:
            estimated_warmup = 15 + params_m * 0.04
            base_speedup = 1.60

        # Estimate time per step (ms) and time saved
        naive_ms = 50 + params_m * 1.0
        compiled_ms = naive_ms / base_speedup
        time_saved_per_step_s = (naive_ms - compiled_ms) / 1000

        if time_saved_per_step_s > 0:
            breakeven = estimated_warmup / time_saved_per_step_s
        else:
            breakeven = float("inf")

        # Return threshold with 1.2x safety margin
        threshold = breakeven * 1.2

        # Optional adjustment based on expected work per step.
        # This uses only pre-run information (batch size, grad accum, seq length).
        generic_scale = 1.0
        fast_scale = 1.0
        if batch_size is not None or grad_accum is not None or max_seq_length is not None:
            try:
                bs = int(batch_size) if batch_size is not None else 2
                ga = int(grad_accum) if grad_accum is not None else 4
                seq = int(max_seq_length) if max_seq_length is not None else 512
            except Exception:
                bs, ga, seq = 2, 4, 512

            bs = max(1, bs)
            ga = max(1, ga)
            # Guard against unbounded tokenizer.model_max_length
            seq = max(64, min(seq, 8192))

            ref_bs, ref_ga, ref_seq = 2, 4, 512

            # Generic path: lighter scaling, less conservative than params-only.
            ga_scale = (ref_ga / ga) ** 1.0
            bs_seq_scale = ((ref_bs * ref_seq) / (bs * seq)) ** 0.15
            generic_scale = 0.35 * ga_scale * bs_seq_scale
            generic_scale = max(0.05, min(generic_scale, 5.0))

            # Fast encoder path: stronger scaling based on observed behavior.
            fast_ga_scale = (ref_ga / ga) ** 1.5
            fast_bs_seq_scale = ((ref_bs * ref_seq) / (bs * seq)) ** 0.25
            fast_scale = 0.2 * fast_ga_scale * fast_bs_seq_scale
            fast_scale = max(0.05, min(fast_scale, 5.0))

        # Conservative safety factors: generic is less conservative than fast.
        generic_threshold = threshold * generic_scale * 1.25

        is_fast_type = (
            isinstance(model_type, str)
            and model_type in FastSentenceTransformer.ENCODER_MODEL_TYPES
        )
        if is_fast_type:
            fast_threshold = threshold * fast_scale * 1.5
            # Prefer the smaller (less conservative) of the two estimates.
            final_threshold = min(generic_threshold, fast_threshold)
        else:
            final_threshold = generic_threshold

        # Reduce mpnet overestimation slightly.
        if model_type == "mpnet":
            final_threshold *= 0.7

        # Lower bound to avoid compiling on extremely short runs.
        return int(max(20, final_threshold))

    @staticmethod
    def _apply_torch_compile(model, mode = "default"):
        """Apply torch.compile to a SentenceTransformer model (with an
        accelerate unwrap_model bug workaround)."""
        if hasattr(model, "__getitem__"):
            inner_model = model[0].auto_model
            compiled = torch.compile(inner_model, mode = mode)
            if isinstance(getattr(type(model[0]), "auto_model", None), property):
                model[0].model = compiled
            else:
                model[0].auto_model = compiled
            # Fix for accelerate unwrap_model bug:
            # When SentenceTransformer contains a compiled inner model,
            # accelerate checks has_compiled_regions() which returns True,
            # then tries to access model.__dict__["_orig_mod"] which fails.
            # This workaround sets _orig_mod to satisfy accelerate.
            model.__dict__["_orig_mod"] = model
        else:
            model = torch.compile(model, mode = mode)
        return model

    @staticmethod
    def from_pretrained(
        model_name,
        max_seq_length = None,
        dtype = None,
        load_in_4bit = False,  # Changed default: 4-bit is slow for encoders
        load_in_8bit = False,
        load_in_16bit = True,  # Changed default: 16-bit is optimal for encoders
        full_finetuning = False,
        token = None,
        device_map = "sequential",
        rope_scaling = None,
        fix_tokenizer = True,
        trust_remote_code = False,
        use_gradient_checkpointing = False,  # Changed default: conflicts with torch.compile
        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,
        unsloth_tiled_mlp = False,
        pooling_mode = "mean",
        for_inference = False,
        **kwargs,
    ):
        try:
            from sentence_transformers import SentenceTransformer
            from sentence_transformers.models import Transformer, Pooling, Normalize
        except ImportError:
            raise ImportError(
                "Unsloth: To use `FastSentenceTransformer`, you must install `sentence-transformers`.\n"
                "Run `pip install sentence-transformers` to install it."
            )

        # Validate the load modes BEFORE the prefetch so a bad config fails without downloading weights.
        # Guard on not for_inference: that branch below never used these flags.
        if not for_inference:
            # sanity check, thanks Etherl:
            if full_finetuning and (load_in_4bit or load_in_8bit):
                print(
                    "Unsloth: You selected full finetuning support, but 4bit / 8bit is enabled - disabling LoRA / QLoRA."
                )
                load_in_4bit = False
                load_in_8bit = False
                load_in_fp8 = False
                load_in_16bit = False

            if int(load_in_4bit) + int(load_in_8bit) + int(load_in_16bit) >= 2:
                raise RuntimeError(
                    "Unsloth: Can only load in 4bit or 8bit or 16bit, not a combination!\n"
                    "Also, we by default set `load_in_16bit = True`.\n"
                    "If you want 4bit LoRA finetuning, set `load_in_16bit = False` and `load_in_4bit = True`\n"
                    "If you want 8bit finetuning, set both `load_in_16bit = False` and `load_in_8bit = True`"
                )

        # Prefetch so the ST load below is a cache hit. weights_at_root stays False (ST component
        # weights live in per-module subfolders). Resolve the same cache the load uses: HF cache_dir,
        # else cache_folder, else SENTENCE_TRANSFORMERS_HOME, else default -- a wrong cache misses the warm.
        _st_prefetched = maybe_prefetch_hf_snapshot(
            model_name,
            token = token,
            revision = revision,
            cache_dir = kwargs.get("cache_dir")
            or kwargs.get("cache_folder")
            or os.environ.get("SENTENCE_TRANSFORMERS_HOME"),
            local_files_only = kwargs.get("local_files_only", False),
            # Forward force_download so the refresh happens in the killable child, then clear it so the
            # in-process ST load reuses the warm cache instead of re-downloading over unguarded Xet.
            force_download = kwargs.get("force_download", False),
        )
        if _st_prefetched and kwargs.get("force_download", False):
            kwargs["force_download"] = False

        # if for_inference == True, skip Unsloth optimizations to avoid torch compile issues
        if for_inference:
            st_device = device_map
            if isinstance(st_device, dict) or (
                isinstance(st_device, str) and st_device in ["auto", "sequential"]
            ):
                st_device = None

            # Propagate dtype to model_kwargs (else inference defaults to float32)
            model_kwargs = kwargs.get("model_kwargs", {})
            model_kwargs["dtype"] = dtype if dtype is not None else "auto"

            st_kwargs = {
                "device": st_device,
                "trust_remote_code": trust_remote_code,
                "token": token,
                "revision": revision,
                "model_kwargs": model_kwargs,
            }

            known_keys = [
                "cache_folder",
                "truncate_dim",
                "tokenizer_kwargs",
                "config_kwargs",
            ]
            for k in known_keys:
                if k in kwargs:
                    st_kwargs[k] = kwargs[k]

            # ST takes cache_folder, not cache_dir: map cache_dir onto it so this load hits the warm
            # (None lets ST honor SENTENCE_TRANSFORMERS_HOME, matching the prefetch).
            _st_cache = kwargs.get("cache_dir") or kwargs.get("cache_folder")
            if _st_cache is not None:
                st_kwargs["cache_folder"] = _st_cache

            st_model = SentenceTransformer(model_name, **st_kwargs)
            return st_model

        # Load-mode validation already ran before the prefetch above.
        if "auto_model" not in kwargs:
            kwargs["auto_model"] = AutoModel

        transformers4 = Version(transformers.__version__).major < 5
        model_type = ""
        config = None
        try:
            config = AutoConfig.from_pretrained(
                model_name, token = token, trust_remote_code = trust_remote_code
            )
            model_type = getattr(config, "model_type", "")
        except:
            pass

        # Fast encoder path: Use native torch.compile for encoder models (6x speedup)
        # This bypasses Unsloth's auto-compiler which adds @torch.compiler.disable decorators
        # that interfere with torch.compile and cause runtime errors for encoder models.
        # NOTE: The old Unsloth path is BROKEN for encoder models with torch 2.9+ due to
        # conflicting @torch.compile and @torch.compiler.disable decorators.
        # Set UNSLOTH_COMPILE_DISABLE=1 to disable torch.compile and use the old path.
        is_encoder_model = model_type.lower() in FastSentenceTransformer.ENCODER_MODEL_TYPES
        use_fast_encoder = os.environ.get("UNSLOTH_COMPILE_DISABLE", "0") != "1"
        if use_fast_encoder and is_encoder_model:
            # torch.compile mode: "default" is safest for PEFT/LoRA training
            # Note: "reduce-overhead" uses CUDA Graphs which is incompatible with PEFT
            compile_mode = "default"

            # Determine dtype - handle float16 machines that don't support bfloat16
            if dtype is None:
                if load_in_16bit:
                    dtype = torch.float16 if not SUPPORTS_BFLOAT16 else torch.bfloat16
                else:
                    dtype = torch.float32
            elif dtype == torch.bfloat16 and not SUPPORTS_BFLOAT16:
                print("Unsloth: Device does not support bfloat16. Using float16 instead.")
                dtype = torch.float16

            # Determine device
            st_device = device_map
            if isinstance(st_device, dict) or (
                isinstance(st_device, str) and st_device in ["auto", "sequential"]
            ):
                st_device = "cuda"

            # Build model_kwargs for SentenceTransformer
            model_kwargs = {"torch_dtype": dtype}

            encoder_attn_impl = resolve_encoder_attention_implementation(
                kwargs.get("auto_model", AutoModel),
                config,
                model_type = model_type,
                disable_sdpa_model_names = DISABLE_SDPA_MODEL_NAMES,
            )
            supports_sdpa = encoder_attn_impl == "sdpa"
            if encoder_attn_impl is not None:
                model_kwargs["attn_implementation"] = encoder_attn_impl

            # Print optimization status
            sdpa_str = " + SDPA" if supports_sdpa else ""
            if load_in_4bit:
                print(
                    f"Unsloth: Using fast encoder path for {model_type} with 4-bit quantization{sdpa_str}"
                )
            else:
                print(
                    f"Unsloth: Using fast encoder path for {model_type} (torch.compile{sdpa_str})"
                )

            # Handle 4-bit quantization via BitsAndBytesConfig
            if load_in_4bit:
                from transformers import BitsAndBytesConfig

                bnb_config = BitsAndBytesConfig(
                    load_in_4bit = True,
                    bnb_4bit_compute_dtype = dtype,
                    bnb_4bit_quant_type = "nf4",
                    bnb_4bit_use_double_quant = True,
                )
                model_kwargs["quantization_config"] = bnb_config
                # When using quantization, device must be handled by accelerate
                st_device = None

            # Handle gradient checkpointing - warn user it conflicts with torch.compile
            _use_gc = use_gradient_checkpointing
            if _use_gc and _use_gc != False:
                print("Unsloth Warning: Gradient checkpointing is incompatible with torch.compile.")
                print("Disabling torch.compile to enable gradient checkpointing.")
                compile_mode = None  # Disable compilation

                is_mpnet = "mpnet" == model_type.lower()

                if is_mpnet and transformers4:
                    FastSentenceTransformer._patch_mpnet_v4()
                elif is_mpnet:
                    FastSentenceTransformer._patch_mpnet_v5()

            # ST takes cache_folder, not cache_dir: map cache_dir onto it so this load hits the warm
            # (None lets ST honor SENTENCE_TRANSFORMERS_HOME, matching the prefetch).
            st_model = SentenceTransformer(
                model_name,
                device = st_device,
                trust_remote_code = trust_remote_code,
                token = token,
                revision = revision,
                model_kwargs = model_kwargs,
                cache_folder = kwargs.get("cache_dir") or kwargs.get("cache_folder"),
            )

            # Store metadata for get_peft_model
            st_model._unsloth_fast_encoder = True
            st_model._compile_mode = compile_mode
            st_model._dtype = dtype
            st_model._load_in_4bit = load_in_4bit
            st_model.no_modules = False
            FastSentenceTransformer._patch_transformer_module_save_config(
                st_model[0], getattr(st_model[0].auto_model, "config", None)
            )

            # Add save methods
            def _save_pretrained_merged(self, save_directory, **save_kwargs):
                self.save_pretrained(save_directory)
                tokenizer = save_kwargs.pop("tokenizer", self.tokenizer)
                if hasattr(self[0], "auto_model"):
                    inner = self[0].auto_model
                    # Handle compiled model
                    if hasattr(inner, "_orig_mod"):
                        inner = inner._orig_mod
                    if hasattr(inner, "merge_and_unload"):
                        merged = inner.merge_and_unload()
                        merged.save_pretrained(save_directory)
                    elif hasattr(inner, "save_pretrained"):
                        inner.save_pretrained(save_directory)
                if tokenizer is not None:
                    tokenizer.save_pretrained(save_directory)
                FastSentenceTransformer._add_unsloth_branding(save_directory)

            st_model.save_pretrained_merged = types.MethodType(_save_pretrained_merged, st_model)

            st_model.save_pretrained_torchao = types.MethodType(_save_pretrained_torchao, st_model)

            st_model.save_pretrained_gguf = types.MethodType(_save_pretrained_gguf, st_model)

            st_model.push_to_hub_gguf = types.MethodType(_push_to_hub_gguf, st_model)

            def _push_to_hub_merged(self, repo_id, **push_kwargs):
                hub_token = push_kwargs.get("token", None) or get_token()
                if hub_token is None:
                    raise ValueError("No HF token provided")
                api = HfApi(token = hub_token)
                try:
                    api.create_repo(
                        repo_id = repo_id,
                        private = push_kwargs.get("private"),
                        exist_ok = True,
                        repo_type = "model",
                    )
                except:
                    pass
                FastSentenceTransformer._add_unsloth_tags(repo_id, hub_token)
                with tempfile.TemporaryDirectory() as temp_dir:
                    self.save_pretrained_merged(temp_dir, **push_kwargs)
                    api.upload_folder(
                        folder_path = temp_dir,
                        repo_id = repo_id,
                        commit_message = push_kwargs.get("commit_message", "Upload model"),
                    )
                print(f"Unsloth: Pushed to https://huggingface.co/{repo_id}")

            st_model.push_to_hub_merged = types.MethodType(_push_to_hub_merged, st_model)

            return st_model

        # Warn if using 4-bit with encoder (slow due to dequantization overhead)
        if is_encoder_model and load_in_4bit:
            print("Unsloth Warning: 4-bit quantization adds ~2.3x overhead for encoder models.")
            print("Consider using load_in_16bit=True for better performance.")

        # check if the model supports add_pooling_layer
        if "add_pooling_layer" not in kwargs:
            supported = FastSentenceTransformer._has_add_pooling_layer(
                config, kwargs.get("auto_model", AutoModel)
            )
            if supported:
                kwargs["add_pooling_layer"] = False

        # fp8 is not supported, force it off
        fp8 = kwargs.pop("load_in_fp8", None)
        if fp8:
            logging.info("Unsloth: Disabling fp8 for model")
        load_in_fp8 = False

        # this is a fix for Snowflake/snowflake-arctic-embed-l-v2.0
        # it has pooler weights which we don't care about for training,
        # however unsloth throws an exception if "UNSLOTH_WARN_UNINITIALIZED" == 1 and it sees unused weights
        old_environ = os.environ.get("UNSLOTH_WARN_UNINITIALIZED", "1")
        os.environ["UNSLOTH_WARN_UNINITIALIZED"] = "0"

        is_distilbert = "distilbert" == model_type.lower()
        is_mpnet = "mpnet" == model_type.lower()

        if is_distilbert and transformers4:
            FastSentenceTransformer._patch_distilbert_v4()
        elif is_distilbert:
            FastSentenceTransformer._patch_distilbert_v5()
        elif is_mpnet and transformers4:
            FastSentenceTransformer._patch_mpnet_v4()
        elif is_mpnet:
            FastSentenceTransformer._patch_mpnet_v5()

        # No modules.json -> force 16-bit: saving is custom for these models and
        # 4-bit would need dequant in save_pretrained_merged, not worth it.
        # Resolve the warmed cache: hf_hub_download ignores SENTENCE_TRANSFORMERS_HOME, so pass it as cache_dir.
        has_modules_json = (
            FastSentenceTransformer._module_path(
                model_name,
                token,
                cache_dir = kwargs.get("cache_dir")
                or kwargs.get("cache_folder")
                or os.environ.get("SENTENCE_TRANSFORMERS_HOME"),
                revision = revision,
            )
            is not None
        )

        if not has_modules_json and load_in_4bit:
            print(
                "Unsloth: No modules.json found. This is not a sentence-transformers model.\n"
                "Forcing 16-bit loading to simplify merged model saving."
            )
            load_in_4bit = False
            load_in_16bit = True

        # The fallback FastModel load reads HF cache_dir, not ST's cache_folder/SENTENCE_TRANSFORMERS_HOME.
        # Point it at the warmed cache, but only when no explicit cache_dir was passed (which wins).
        _st_cache_dir = kwargs.get("cache_folder") or os.environ.get("SENTENCE_TRANSFORMERS_HOME")
        if _st_cache_dir is not None and "cache_dir" not in kwargs:
            kwargs["cache_dir"] = _st_cache_dir

        try:
            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,
            )
        finally:
            os.environ["UNSLOTH_WARN_UNINITIALIZED"] = old_environ

        from sentence_transformers import SentenceTransformer

        modules, no_modules = FastSentenceTransformer._load_modules(
            model_name,
            token,
            model,
            tokenizer,
            max_seq_length,
            pooling_mode,
            trust_remote_code = trust_remote_code,
            # Same resolved cache as above so the fallback module loads hit the warm, not Xet.
            cache_dir = kwargs.get("cache_dir")
            or kwargs.get("cache_folder")
            or os.environ.get("SENTENCE_TRANSFORMERS_HOME"),
            # Same revision as the weight load so modules hit the warm (None = default branch).
            revision = revision,
        )

        st_device = device_map
        if isinstance(st_device, dict) or (
            isinstance(st_device, str) and st_device in ["auto", "sequential"]
        ):
            st_device = None

        st_model = SentenceTransformer(modules = modules, device = st_device)
        st_model.no_modules = no_modules

        def _save_pretrained_merged(self, save_directory, **kwargs):
            # check which adapter files exist before save_pretrained
            adapter_files = ["adapter_model.safetensors", "adapter_config.json"]
            existing_before = {
                f for f in adapter_files if os.path.exists(os.path.join(save_directory, f))
            }

            # sentence-transformers config and modules only get saved if we call save_pretrained
            self.save_pretrained(save_directory)

            # remove LoRA adapters only if they were created by save_pretrained (not pre-existing)
            for file in adapter_files:
                if file not in existing_before:
                    try:
                        os.remove(os.path.join(save_directory, file))
                    except:
                        pass

            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...")
                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",
                ]
                for k in unsloth_args:
                    safe_kwargs.pop(k, None)

                merged_model = self[0].auto_model.merge_and_unload()
                merged_model.save_pretrained(save_directory, **safe_kwargs)
                if tokenizer is not None:
                    tokenizer.save_pretrained(save_directory)
            else:
                self[0].auto_model.save_pretrained_merged(
                    save_directory, tokenizer = tokenizer, **kwargs
                )

            # add Unsloth branding to the generated README
            try:
                FastSentenceTransformer._add_unsloth_branding(save_directory)
            except Exception as e:
                print(f"Unsloth Warning: Failed to add branding to README: {e}")

        st_model.save_pretrained_merged = types.MethodType(_save_pretrained_merged, st_model)

        st_model.save_pretrained_torchao = types.MethodType(_save_pretrained_torchao, st_model)

        st_model.save_pretrained_gguf = types.MethodType(_save_pretrained_gguf, st_model)

        st_model.push_to_hub_gguf = types.MethodType(_push_to_hub_gguf, st_model)

        def _push_to_hub_merged(self, repo_id, **kwargs):
            token = kwargs.get("token", None) or get_token()
            if token is None:
                raise ValueError(
                    "No HF token provided. Please provide a token or login with `hf auth login`"
                )
            private = kwargs.get("private", None)
            commit_message = kwargs.get("commit_message", "Upload model")

            from huggingface_hub import HfApi

            api = HfApi(token = token)
            try:
                api.create_repo(
                    repo_id = repo_id,
                    private = private,
                    exist_ok = True,
                    repo_type = "model",
                )
            except:
                pass

            # order doesn't seem to matter for this after repo creation...
            FastSentenceTransformer._add_unsloth_tags(repo_id, token)

            with tempfile.TemporaryDirectory() as temp_dir:
                self.save_pretrained_merged(temp_dir, **kwargs)
                api.upload_folder(
                    folder_path = temp_dir,
                    repo_id = repo_id,
                    commit_message = commit_message,
                )
            print(f"Unsloth: Successfully pushed merged model to https://huggingface.co/{repo_id}")

        st_model.push_to_hub_merged = types.MethodType(_push_to_hub_merged, st_model)
        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 = False,  # Changed default: conflicts with torch.compile
        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
        from peft import LoraConfig, get_peft_model as peft_get_peft_model

        if "task_type" not in kwargs:
            kwargs["task_type"] = "FEATURE_EXTRACTION"
            print("Setting task_type to FEATURE_EXTRACTION")

        if isinstance(model, SentenceTransformer):
            # Check if this is a fast encoder model (uses torch.compile instead of Unsloth patching)
            is_fast_encoder = getattr(model, "_unsloth_fast_encoder", False)

            if is_fast_encoder:
                # Fast encoder path: Use native PEFT + torch.compile (6x speedup)
                transformer_module = model[0]
                inner_model = transformer_module.auto_model

                # Check if model is quantized (4-bit/8-bit)
                is_quantized = (
                    getattr(inner_model, "is_quantized", False)
                    or getattr(inner_model.config, "quantization_config", None) is not None
                )

                # Track if gradient checkpointing was actually enabled
                gc_enabled = False

                # this is needed when from_pretrained was called without gradient
                # checkpointing but get_peft_model requests it
                if use_gradient_checkpointing and use_gradient_checkpointing != False:
                    import transformers
                    from packaging.version import Version

                    transformers4 = Version(transformers.__version__).major < 5
                    model_type = getattr(inner_model.config, "model_type", "").lower()

                    if model_type == "mpnet" and transformers4:
                        FastSentenceTransformer._patch_mpnet_v4()
                    elif model_type == "mpnet":
                        FastSentenceTransformer._patch_mpnet_v5()

                # Prepare for k-bit training if quantized
                if is_quantized:
                    from ._utils import prepare_model_for_kbit_training
                    _gc_for_kbit = (
                        use_gradient_checkpointing if use_gradient_checkpointing else False
                    )
                    try:
                        inner_model = prepare_model_for_kbit_training(
                            inner_model,
                            use_gradient_checkpointing = _gc_for_kbit,
                        )
                        print("Unsloth: Prepared quantized model for k-bit training")
                        gc_enabled = bool(_gc_for_kbit)
                    except ValueError as e:
                        if "does not support gradient checkpointing" in str(e):
                            # Model doesn't support gradient checkpointing, disable it
                            print(
                                f"Unsloth Warning: {inner_model.__class__.__name__} does not support gradient checkpointing. Skipping."
                            )
                            inner_model = prepare_model_for_kbit_training(
                                inner_model,
                                use_gradient_checkpointing = False,
                            )
                            print(
                                "Unsloth: Prepared quantized model for k-bit training (without gradient checkpointing)"
                            )
                        else:
                            raise

                # Enable gradient checkpointing if requested (only for non-quantized, since prepare_model handles it)
                elif use_gradient_checkpointing and use_gradient_checkpointing != False:
                    if hasattr(inner_model, "gradient_checkpointing_enable"):
                        try:
                            inner_model.gradient_checkpointing_enable()
                            print("Unsloth: Enabled gradient checkpointing")
                            gc_enabled = True
                        except ValueError as e:
                            if "does not support gradient checkpointing" in str(e):
                                print(
                                    f"Unsloth Warning: {inner_model.__class__.__name__} does not support gradient checkpointing. Skipping."
                                )

                # Create LoRA config
                lora_config = LoraConfig(
                    r = r,
                    lora_alpha = lora_alpha,
                    target_modules = target_modules,
                    lora_dropout = lora_dropout,
                    bias = bias,
                    task_type = kwargs.get("task_type", "FEATURE_EXTRACTION"),
                )

                # Apply PEFT directly (not through FastModel)
                peft_model = peft_get_peft_model(inner_model, lora_config)

                # Apply QAT if specified
                qat_scheme = kwargs.get("qat_scheme", None)
                if qat_scheme is not None:
                    from ._utils import _prepare_model_for_qat
                    peft_model = _prepare_model_for_qat(peft_model, qat_scheme)

                # Determine compile mode (only if not using gradient checkpointing)
                compile_mode = getattr(model, "_compile_mode", "default")
                # Re-enable torch.compile if gradient checkpointing was requested but couldn't be enabled
                if compile_mode is None and not gc_enabled:
                    compile_mode = "default"
                    print(
                        "Unsloth: Re-enabling torch.compile since gradient checkpointing is not supported"
                    )

                # Re-assign the peft model back to the transformer module.
                # On sentence-transformers >=5.4 `auto_model` is a read-only property
                # backed by `self.model`, so write to the backing attribute there.
                if isinstance(getattr(type(transformer_module), "auto_model", None), property):
                    transformer_module.model = peft_model
                else:
                    transformer_module.auto_model = peft_model
                FastSentenceTransformer._patch_transformer_module_save_config(
                    transformer_module, getattr(inner_model, "config", None)
                )

                # Store compile info for auto-compile at trainer time
                # torch.compile is deferred until training starts so we can check max_steps
                if compile_mode is not None:
                    model._compile_mode = compile_mode
                    model._compile_threshold = FastSentenceTransformer._estimate_compile_threshold(
                        model
                    )
                    # Flag to indicate compile has not been applied yet
                    model._compile_pending = True
                    print(
                        f"Unsloth: torch.compile will be applied automatically if max_steps > {model._compile_threshold}"
                    )
                else:
                    model._compile_mode = None
                    model._compile_pending = False
                    print("Unsloth: torch.compile disabled (gradient checkpointing enabled)")

                return model

            # Original path for non-fast-encoder models
            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.
            # On sentence-transformers >=5.4 `auto_model` is a read-only property
            # backed by `self.model`, so write to the backing attribute there.
            if isinstance(getattr(type(transformer_module), "auto_model", None), property):
                transformer_module.model = peft_model
            else:
                transformer_module.auto_model = peft_model
            FastSentenceTransformer._patch_transformer_module_save_config(
                transformer_module, getattr(inner_model, "config", None)
            )
            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,
            )


def _patch_sentence_transformer_trainer():
    """
    Patch SentenceTransformerTrainer to automatically apply torch.compile
    when training steps exceed the breakeven threshold.

    This is called automatically when this module is imported.
    """
    try:
        from sentence_transformers import SentenceTransformerTrainer
    except ImportError:
        return  # sentence_transformers not installed

    if getattr(SentenceTransformerTrainer, "_unsloth_auto_compile_patched", False):
        return  # Already patched

    from functools import wraps

    _original_init = SentenceTransformerTrainer.__init__

    @wraps(_original_init)
    def _patched_init(self, *args, **kwargs):
        # Extract model and training_args
        model = kwargs.get("model") or (args[0] if args else None)
        training_args = kwargs.get("args") or (args[1] if len(args) > 1 else None)

        # Check if model has pending compile
        if (
            model is not None
            and training_args is not None
            and getattr(model, "_compile_pending", False)
        ):
            max_steps = getattr(training_args, "max_steps", -1)
            compile_mode = getattr(model, "_compile_mode", "default")

            # Re-estimate threshold now that training args are available
            batch_size = getattr(training_args, "per_device_train_batch_size", None)
            grad_accum = getattr(training_args, "gradient_accumulation_steps", None)
            max_seq_length = getattr(model, "max_seq_length", None)
            if max_seq_length is None and hasattr(model, "__getitem__"):
                try:
                    max_seq_length = getattr(model[0], "max_seq_length", None)
                except Exception:
                    max_seq_length = None
            if max_seq_length is None:
                tokenizer = getattr(model, "tokenizer", None)
                max_seq_length = (
                    getattr(tokenizer, "model_max_length", None) if tokenizer is not None else None
                )

            threshold = FastSentenceTransformer._estimate_compile_threshold(
                model,
                batch_size = batch_size,
                grad_accum = grad_accum,
                max_seq_length = max_seq_length,
            )
            model._compile_threshold = threshold

            if max_steps > 0 and max_steps >= threshold:
                print(f"Unsloth: Auto-compiling model ({max_steps} steps >= {threshold} threshold)")
                FastSentenceTransformer._apply_torch_compile(model, mode = compile_mode)
                model._compile_pending = False
            elif max_steps > 0:
                print(
                    f"Unsloth: Skipping torch.compile ({max_steps} steps < {threshold} threshold)"
                )
                model._compile_pending = False

        # Call original __init__
        _original_init(self, *args, **kwargs)

        # Disable mixed precision when FORCE_FLOAT32 is active (matches rl.py behavior)
        if os.environ.get("UNSLOTH_FORCE_FLOAT32", "0") == "1":
            if hasattr(self, "args") and self.args is not None:
                if self.args.fp16 or self.args.bf16:
                    print(
                        "Unsloth: Switching to float32 training since model cannot work with float16"
                    )
                    self.args.fp16 = False
                    self.args.bf16 = False
                    if hasattr(self.args, "bf16_full_eval"):
                        self.args.bf16_full_eval = False
                    if hasattr(self.args, "fp16_full_eval"):
                        self.args.fp16_full_eval = False

    SentenceTransformerTrainer.__init__ = _patched_init
    SentenceTransformerTrainer._unsloth_auto_compile_patched = True


def _patch_st_trainer_load_from_checkpoint():
    try:
        from sentence_transformers import SentenceTransformerTrainer
    except ImportError:
        return
    if getattr(SentenceTransformerTrainer, "_unsloth_load_from_checkpoint_patched", False):
        return
    if not hasattr(SentenceTransformerTrainer, "_load_from_checkpoint"):
        return

    _original = SentenceTransformerTrainer._load_from_checkpoint

    def _unsloth_load_from_checkpoint(self, checkpoint_path):
        try:
            from peft import PeftModel, load_peft_weights, set_peft_model_state_dict
        except ImportError:
            return _original(self, checkpoint_path)

        try:
            mod0 = self.model[0]
        except (IndexError, TypeError):
            return _original(self, checkpoint_path)

        if isinstance(getattr(type(mod0), "auto_model", None), property):
            inner = getattr(mod0, "model", None)
        else:
            inner = getattr(mod0, "auto_model", None)
        inner = getattr(inner, "_orig_mod", inner)

        if not isinstance(inner, PeftModel):
            return _original(self, checkpoint_path)
        if not getattr(mod0, "_unsloth_st_managed", False):
            return _original(self, checkpoint_path)

        if not any(
            os.path.isfile(os.path.join(checkpoint_path, fn))
            for fn in ("adapter_model.safetensors", "adapter_model.bin")
        ):
            return _original(self, checkpoint_path)

        adapter_name = getattr(inner, "active_adapter", None)
        if adapter_name is None and callable(getattr(inner, "active_adapters", None)):
            adapter_name = inner.active_adapters()
        if isinstance(adapter_name, (list, tuple, set)):
            if len(adapter_name) != 1:
                raise RuntimeError("Unsloth: Cannot resume multiple active PEFT adapters.")
            adapter_name = next(iter(adapter_name))
        adapter_name = adapter_name or "default"
        if adapter_name not in getattr(inner, "peft_config", {}):
            raise RuntimeError(f"Unsloth: PEFT adapter {adapter_name!r} is not loaded.")

        load_result = set_peft_model_state_dict(
            inner, load_peft_weights(checkpoint_path), adapter_name = adapter_name
        )
        unexpected = getattr(load_result, "unexpected_keys", []) or []
        missing = [
            x
            for x in (getattr(load_result, "missing_keys", []) or [])
            if f".{adapter_name}." in x or x.endswith(f".{adapter_name}")
        ]
        if unexpected or missing:
            raise RuntimeError(
                "Unsloth: PEFT checkpoint does not match the active adapter "
                f"(missing={missing[:8]}, unexpected={unexpected[:8]})."
            )

        modules_json = os.path.join(checkpoint_path, "modules.json")
        if not os.path.isfile(modules_json):
            raise RuntimeError("Unsloth: PEFT checkpoint is missing modules.json.")
        try:
            with open(modules_json, "r") as f:
                module_configs = json.load(f)
        except Exception as e:
            raise RuntimeError("Unsloth: Cannot parse checkpoint modules.json.") from e

        root = os.path.abspath(os.fspath(checkpoint_path))
        restored = set()
        for entry in module_configs:
            idx = int(entry.get("idx", -1))
            if idx == 0:
                continue
            if idx < 0 or idx >= len(self.model):
                raise RuntimeError(f"Unsloth: Bad module index in modules.json: {idx}.")
            module = self.model[idx]
            module_cls = type(module)
            saved_type = entry.get("type", "")
            if saved_type and not saved_type.endswith(f".{module_cls.__name__}"):
                raise RuntimeError(f"Unsloth: Checkpoint module {idx} type mismatch.")
            module_path = entry.get("path")
            module_dir = os.path.abspath(os.path.join(root, os.fspath(module_path or "")))
            try:
                inside_root = os.path.commonpath([root, module_dir]) == root
            except ValueError:
                inside_root = False
            if not module_path or not inside_root or not os.path.isdir(module_dir):
                raise RuntimeError(f"Unsloth: Bad checkpoint module path for index {idx}.")
            if not hasattr(module_cls, "load"):
                raise RuntimeError(f"Unsloth: Module {idx} cannot be reloaded.")
            fresh = module_cls.load(module_dir)
            if not isinstance(fresh, module_cls):
                raise RuntimeError(f"Unsloth: Module {idx} reload returned wrong type.")
            # Parameterless modules (Pooling, Normalize) make
            # next(module.parameters()) raise StopIteration; route through
            # the SentenceTransformer's device property instead.
            try:
                fresh.to(self.model.device)
            except AttributeError:
                pass
            self.model[idx] = fresh
            restored.add(idx)
        missing_idx = sorted(set(range(1, len(self.model))) - restored)
        if missing_idx:
            raise RuntimeError(
                f"Unsloth: Checkpoint modules.json is incomplete (missing idx={missing_idx[:8]})."
            )

    SentenceTransformerTrainer._load_from_checkpoint = _unsloth_load_from_checkpoint
    SentenceTransformerTrainer._unsloth_load_from_checkpoint_patched = True


_patch_sentence_transformer_trainer()
_patch_st_trainer_load_from_checkpoint()