from __future__ import annotations import importlib import inspect import json import os import shutil import sys import tempfile import traceback from abc import ABC, abstractmethod from collections import OrderedDict from multiprocessing import Queue from pathlib import Path from typing import Any, Literal import numpy as np import torch import torch.multiprocessing as mp import transformers from huggingface_hub import CardData, HfApi from packaging import version from torch import Tensor, nn from transformers import PreTrainedModel, is_datasets_available, is_torch_npu_available from transformers.dynamic_module_utils import get_class_from_dynamic_module, get_relative_import_files from transformers.utils import logging as transformers_logging from sentence_transformers import __version__ from sentence_transformers.base.evaluation import BaseEvaluator from sentence_transformers.base.modality import format_modality, infer_batch_modality from sentence_transformers.base.modality_types import Modality, PairInput, SingleInput from sentence_transformers.base.model_card import BaseModelCardData, generate_model_card from sentence_transformers.base.modules import Module, Router, Transformer from sentence_transformers.base.peft_mixin import PeftAdapterMixin from sentence_transformers.util import ( get_device_name, import_from_string, load_dir_path, load_file_path, save_to_hub_args_decorator, ) from sentence_transformers.util.misc import ORIGINAL_TRANSFORMER_MODELS logger = transformers_logging.get_logger(__name__) class BaseModel(nn.Sequential, PeftAdapterMixin, ABC): """ Base class for SentenceTransformer, SparseEncoder, and CrossEncoder models. This class provides common functionality for: - Model loading (from Hub, local paths, or creating new models) - Model saving (to disk and Hub) - Device management - Module architecture (sequential composition) - Configuration management - Tokenizer/processor access All models inherit from nn.Sequential and are composed of a sequence of modules that are called sequentially in the forward pass. """ # The dataclass used to generate the model card when saving the model. model_card_data_class = BaseModelCardData # The default Hugging Face organization to prepend to short model names. default_huggingface_organization: str | None = None # Default prompts to initialize for new model instances. Use None as the value for prompts # that should be filled by the saved model config; empty string "" means intentionally blank. _default_prompts: dict[str, str | None] = {} # The placeholder model ID in model card templates that gets replaced with the actual model ID. _model_card_model_id_placeholder: str = "sentence_transformers_model_id" def __init__( self, model_name_or_path: str | None = None, *, modules: list[nn.Module] | OrderedDict[str, nn.Module] | None = None, device: str | None = None, prompts: dict[str, str] | None = None, default_prompt_name: str | None = None, cache_folder: str | None = None, trust_remote_code: bool = False, revision: str | None = None, local_files_only: bool = False, token: bool | str | None = None, model_kwargs: dict[str, Any] | None = None, processor_kwargs: dict[str, Any] | None = None, config_kwargs: dict[str, Any] | None = None, model_card_data: CardData | None = None, backend: Literal["torch", "onnx", "openvino"] = "torch", ) -> None: """ Initialize a BaseModel instance. Args: model_name_or_path (str, optional): If a filepath on disk, loads the model from that path. Otherwise, tries to download a pre-trained model. If that fails, tries to construct a model from the Hugging Face Hub with that name. Defaults to None. modules (list[nn.Module], optional): A list of torch modules that are called sequentially. Can be used to create custom models from scratch. Defaults to None. device (str, optional): Device (like ``"cuda"``, ``"cpu"``, ``"mps"``, ``"npu"``) that should be used for computation. If None, checks if a GPU can be used. Defaults to None. prompts (dict[str, str], optional): A dictionary with prompts for the model. The key is the prompt name, the value is the prompt text. The prompt text will be prepended before any text during inference. For example: ``{"query": "query: ", "passage": "passage: "}``. If a model has saved prompts, you can override them by passing your own, or pass ``{"query": "", "document": ""}`` to disable them. Defaults to None. default_prompt_name (str, optional): The name of the prompt that should be used by default. If not set, no prompt will be applied. Defaults to None. cache_folder (str, optional): Path to store models. Can also be set by the ``SENTENCE_TRANSFORMERS_HOME`` environment variable. Defaults to None. trust_remote_code (bool, optional): Whether to allow for custom models defined on the Hub in their own modeling files. Only set to ``True`` for repositories you trust and in which you have read the code, as it will execute code present on the Hub on your local machine. Defaults to False. revision (str, optional): The specific model version to use. It can be a branch name, a tag name, or a commit id, for a stored model on Hugging Face. Defaults to None. local_files_only (bool, optional): Whether to only look at local files (i.e., do not try to download the model). Defaults to False. token (bool or str, optional): Hugging Face authentication token to download private models. Defaults to None. model_kwargs (dict[str, Any], optional): Keyword arguments passed to the underlying Hugging Face Transformers model via ``AutoModel.from_pretrained``. Particularly useful options include: - ``torch_dtype``: Override the default ``torch.dtype`` and load the model under a specific dtype. Can be ``torch.float16``, ``torch.bfloat16``, ``torch.float32``, or ``"auto"`` to use the dtype from the model's ``config.json``. - ``attn_implementation``: The attention implementation to use. For example ``"eager"``, ``"sdpa"``, or ``"flash_attention_2"``. If you ``pip install kernels``, then ``"flash_attention_2"`` should work without having to install ``flash_attn``. It is frequently the fastest option. Defaults to ``"sdpa"`` when available (torch>=2.1.1). - ``device_map``: Device map for model parallelism, e.g. ``"auto"``. - ``provider``: For ``backend="onnx"``, the ONNX execution provider (e.g. ``"CUDAExecutionProvider"``). - ``file_name``: For ``backend="onnx"`` or ``"openvino"``, the filename to load (e.g. for optimized or quantized models). - ``export``: For ``backend="onnx"`` or ``"openvino"``, whether to export the model to the backend format. Also set automatically if the exported file doesn't exist. See the `PreTrainedModel.from_pretrained `_ documentation for more details. Defaults to None. processor_kwargs (dict[str, Any], optional): Keyword arguments passed to the Hugging Face Transformers processor/tokenizer via ``AutoProcessor.from_pretrained``. See the `AutoTokenizer.from_pretrained `_ documentation for more details. Defaults to None. config_kwargs (dict[str, Any], optional): Keyword arguments passed to the Hugging Face Transformers config via ``AutoConfig.from_pretrained``. See the `AutoConfig.from_pretrained `_ documentation for more details. Defaults to None. model_card_data (CardData, optional): A model card data object that contains information about the model. Used to generate a model card when saving the model. If not set, a default model card data object is created. Defaults to None. backend (str, optional): The backend to use for inference. Can be ``"torch"`` (default), ``"onnx"``, or ``"openvino"``. Defaults to ``"torch"``. """ default_prompts = dict(self._default_prompts) if prompts: default_prompts.update(prompts) self.prompts = default_prompts self.default_prompt_name = default_prompt_name self.trust_remote_code = trust_remote_code self.model_card_data = model_card_data or self.model_card_data_class(local_files_only=local_files_only) self.module_kwargs = None self._model_card_vars = {} self._model_card_text = None self.model_type = self.__class__.__name__ self.backend = backend if cache_folder is None: cache_folder = os.getenv("SENTENCE_TRANSFORMERS_HOME") # Determine device if device is None: device = get_device_name() logger.info(f"No device provided, using {device}") if device == "hpu" and importlib.util.find_spec("optimum") is not None: from optimum.habana.transformers.modeling_utils import adapt_transformers_to_gaudi adapt_transformers_to_gaudi() # Load model if model_name_or_path and not os.path.exists(model_name_or_path): # Not a local path, load from hub if (os.sep == "\\" and "\\" in model_name_or_path) or model_name_or_path.count("/") > 1: raise FileNotFoundError(f"Path {model_name_or_path} not found") if ( self.default_huggingface_organization is not None and "/" not in model_name_or_path and model_name_or_path.lower() not in ORIGINAL_TRANSFORMER_MODELS ): model_name_or_path = f"{self.default_huggingface_organization}/{model_name_or_path}" if model_name_or_path: modules, self.module_kwargs = self._load_modules( model_name_or_path, token=token, cache_folder=cache_folder, revision=revision, trust_remote_code=trust_remote_code, local_files_only=local_files_only, model_kwargs=model_kwargs, processor_kwargs=processor_kwargs, config_kwargs=config_kwargs, ) if modules is not None and not isinstance(modules, OrderedDict): modules = OrderedDict([(str(idx), module) for idx, module in enumerate(modules)]) if modules is not None and len(modules) == 0: raise ValueError( f"An empty modules list was passed to {type(self).__name__}. " "Please provide at least one module, e.g. a Transformer module." ) super().__init__(modules) # Cast non-input modules to match the first module's dtype for consistency. # The first module (e.g. Transformer) is the dtype source of truth and downstream # modules (Dense, Pooling, etc.) should match it. first_param = next(self[0].parameters(), None) if first_param is not None: first_dtype = first_param.dtype for module in list(self.children())[1:]: module.to(first_dtype) self.to(device) self.is_hpu_graph_enabled = False # Validate prompts after model loading (which may have merged config prompts) self._validate_prompts() # Pass the model to the model card data for later use self.model_card_data.register_model(self) def _validate_prompts(self) -> None: """Validate prompt configuration and log prompt information.""" # Replace any remaining None prompts (not filled by saved config) with empty strings for key, value in self.prompts.items(): if value is None: self.prompts[key] = "" if self.default_prompt_name is not None and self.default_prompt_name not in self.prompts: raise ValueError( f"Default prompt name '{self.default_prompt_name}' not found in the configured prompts " f"dictionary with keys {list(self.prompts.keys())!r}." ) if non_empty_keys := [k for k, v in self.prompts.items() if v != ""]: n = len(non_empty_keys) logger.info(f"Loaded {n} prompt{'s' if n > 1 else ''} with these keys: {non_empty_keys}") if self.default_prompt_name: logger.warning_once( f"Default prompt name is set to '{self.default_prompt_name}'. " f"This prompt will be applied to all inference calls, except if " f"a `prompt` or `prompt_name` parameter is provided." ) def _resolve_prompt(self, prompt: str | None, prompt_name: str | None) -> str | None: """Resolve a prompt from a prompt name or the default prompt name. Args: prompt: An explicit prompt string, or None. prompt_name: A key into ``self.prompts``, or None. Returns: The resolved prompt string, or None if no prompt applies. """ if prompt is None: if prompt_name is not None: try: prompt = self.prompts[prompt_name] except KeyError: raise ValueError( f"Prompt name '{prompt_name}' not found in the configured prompts dictionary with keys {list(self.prompts.keys())!r}." ) elif self.default_prompt_name is not None: prompt = self.prompts.get(self.default_prompt_name, None) elif prompt_name is not None: logger.warning( "Provide either a `prompt`, a `prompt_name`, or neither, but not both. " "Ignoring the `prompt_name` in favor of `prompt`." ) return prompt def get_backend(self) -> Literal["torch", "onnx", "openvino"]: """Return the backend used for inference, which can be one of "torch", "onnx", or "openvino". Returns: str: The backend used for inference. """ return self.backend @property def modalities(self) -> list[Modality]: """Return the list of modalities supported by this model, e.g. ``["text"]`` or ``["text", "image", "message"]``.""" return getattr(self[0], "modalities", ["text"]) def supports(self, modality: Modality) -> bool: """Check if the model supports the given modality. A modality is supported if: 1. It is directly listed in :attr:`modalities` (including tuple modalities that are explicitly listed), or 2. It is a tuple of modalities (e.g. ``("image", "text")``) where each part is individually supported and the model also supports ``"message"`` format, which is used to combine multiple modalities into a single input. Args: modality: A single modality string (e.g. ``"text"``, ``"image"``) or a tuple of modality strings (e.g. ``("image", "text")``). Returns: bool: Whether the model supports the given modality. Example:: >>> from sentence_transformers import SentenceTransformer >>> model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2") >>> model.supports("text") True >>> model.supports("image") False """ supported = self.modalities if modality in supported: return True if isinstance(modality, tuple) and "message" in supported: return all(part in supported for part in modality) return False def get_model_kwargs(self) -> list[str]: """ Get the keyword arguments specific to this model for inference methods like `encode` or `predict`. Example: >>> from sentence_transformers import SentenceTransformer, SparseEncoder >>> SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2").get_model_kwargs() [] >>> SentenceTransformer("jinaai/jina-embeddings-v4", trust_remote_code=True).get_model_kwargs() ['task', 'truncate_dim'] >>> SparseEncoder("opensearch-project/opensearch-neural-sparse-encoding-doc-v3-distill").get_model_kwargs() ['task'] Returns: list[str]: A list of keyword arguments for the forward pass. """ modules = list(self.named_children()) forward_kwargs = set() while modules: module_name, module = modules.pop() if isinstance(module, Router): for route_modules in module.sub_modules.values(): modules.extend(list(route_modules.named_children())) if self.module_kwargs and module_name in self.module_kwargs: forward_kwargs.update(self.module_kwargs[module_name]) if hasattr(module, "forward_kwargs"): forward_kwargs.update(module.forward_kwargs) return list(forward_kwargs) def get_max_seq_length(self) -> int | None: """ .. deprecated:: Use the :attr:`max_seq_length` property instead. Returns the maximal sequence length that the first module of the model accepts. Longer inputs will be truncated. Returns: Optional[int]: The maximal sequence length that the model accepts, or None if it is not defined. """ return self.max_seq_length def _first_module(self) -> torch.nn.Module: """Returns the first module of this sequential embedder""" return self._modules[next(iter(self._modules))] def _last_module(self) -> torch.nn.Module: """Returns the last module of this sequential embedder""" return self._modules[next(reversed(self._modules))] def _can_flatten_inputs(self) -> bool: """Check if the first module (Transformer or Router containing Transformers) supports flattened text-only inputs. For Router models, this returns True only if ALL routes support flattening, since we cannot know at this point which route will be used. """ input_module = self[0] if isinstance(input_module, Transformer): return input_module.can_flatten_inputs if isinstance(input_module, Router): for route in input_module.sub_modules.values(): # Each route is nn.Sequential; the first child is the Transformer first_in_route = next(iter(route.children()), None) if not isinstance(first_in_route, Transformer) or not first_in_route.can_flatten_inputs: return False return bool(input_module.sub_modules) return False @staticmethod def _interleave_sorted_indices(sorted_idx: np.ndarray) -> np.ndarray: """Interleave a largest-to-smallest sorted index array so that each consecutive batch contains a mix of long and short inputs. When text-only inputs are flattened via flash attention, there is no padding. Total tokens per batch equals the sum of actual lengths. Grouping all long inputs together creates peak-memory batches, while grouping all short inputs together under-utilises the GPU. Interleaving (largest, smallest, 2nd largest, 2nd smallest, ...) balances the total token count across batches for more uniform memory usage. """ n = len(sorted_idx) interleaved = np.empty(n, dtype=sorted_idx.dtype) interleaved[0::2] = sorted_idx[: (n + 1) // 2] interleaved[1::2] = sorted_idx[n - 1 : (n - 1) // 2 : -1] return interleaved @staticmethod def _input_length(sample) -> int: """Estimate the "size" of an input sample for length-based batch sorting. The exact value doesn't matter, it's only used to group similarly sized inputs together so that padding waste is minimised within each batch. """ # Strings -> character count (decent proxy for token count). Checked first # because it's the most common input type and avoids heavier checks. if isinstance(sample, str): return len(sample) # Text pair or non-text pair if isinstance(sample, (tuple, list)): if sample and not (isinstance(sample[0], dict) and "role" in sample[0]): return sum(BaseModel._input_length(s) for s in sample) # Dict inputs: audio/video wrappers, messages, multimodal dicts if isinstance(sample, dict): if "array" in sample: return BaseModel._input_length(sample["array"]) if "content" in sample: content = sample["content"] if isinstance(content, str): return len(content) if isinstance(content, list): return sum(len(str(item.get("text", ""))) for item in content if isinstance(item, dict)) return 1 # Multimodal dict: sum of component sizes return sum(BaseModel._input_length(v) for v in sample.values()) # List of message dicts if isinstance(sample, list) and sample and isinstance(sample[0], dict): return sum(BaseModel._input_length(msg) for msg in sample) # Tensors / arrays -> total number of elements if isinstance(sample, (np.ndarray, torch.Tensor)): return sample.nelement() if isinstance(sample, torch.Tensor) else sample.size # PIL Image -> pixel count try: from PIL.Image import Image as PILImage if isinstance(sample, PILImage): return sample.size[0] * sample.size[1] except ImportError: pass return 1 def forward(self, input: dict[str, Tensor], **kwargs) -> dict[str, Tensor]: """Forward pass through all modules in the model.""" for module_name, module in self.named_children(): module_kwargs = {} if isinstance(module, Router): module_kwargs = kwargs else: module_kwarg_keys = [] if self.module_kwargs is not None: module_kwarg_keys = self.module_kwargs.get(module_name, []) module_kwargs = { key: value for key, value in kwargs.items() if key in module_kwarg_keys or (hasattr(module, "forward_kwargs") and key in module.forward_kwargs) } input = module(input, **module_kwargs) return input def preprocess( self, inputs: list[SingleInput | PairInput], prompt: str | None = None, **kwargs, ) -> dict[str, Tensor | Any]: """ Preprocesses the inputs for the model. Args: inputs (list[SingleInput | PairInput]): A list of inputs to be preprocessed. Each input can be a string, dict, tuple, PIL Image, numpy array, torch Tensor, or other supported modality. If a single input is provided, it must be wrapped in a list. prompt (str, optional): A prompt string to prepend to text inputs. Defaults to None. If the model supports the ``message`` modality, the prompt will be added as a system message to the input messages instead of being prepended to text. Returns: dict[str, Tensor | Any]: A dictionary of tensors with the preprocessed inputs. """ if not inputs: return {} # Validate that the inputs match a supported modality. # If "message" is supported, any modality is allowed since the input module # can convert it to message format (e.g. wrapping images in chat messages). modality = None try: modality = infer_batch_modality(inputs, supported_modalities=self.modalities) except (ValueError, TypeError): pass if modality is not None and not self.supports(modality): supported = ", ".join(format_modality(m) for m in self.modalities) message = ( f"Modality '{format_modality(modality)}' is not supported by this {type(self).__name__} model. " f"Supported modalities: {supported}" ) if isinstance(modality, tuple) and all(part in self.modalities for part in modality): message += ( f"\nThis model supports {' and '.join(modality)} individually, " "but not in the same input. Please process each modality separately." ) raise ValueError(message) # Backwards compatibility: fall back to preprocess/tokenize without prompt if the # input module doesn't accept it. Only the main path (preprocess with prompt) will # be supported in the future. try: preprocessed = self[0].preprocess(inputs, prompt=prompt, **kwargs) except TypeError: if prompt and modality == "text": inputs = [(prompt + inp[0],) + inp[1:] if isinstance(inp, tuple) else prompt + inp for inp in inputs] preprocessed = self[0].preprocess(inputs, **kwargs) except AttributeError: if prompt and modality == "text": inputs = [(prompt + inp[0],) + inp[1:] if isinstance(inp, tuple) else prompt + inp for inp in inputs] try: preprocessed = self[0].tokenize(inputs, **kwargs) except TypeError: preprocessed = self[0].tokenize(inputs) return preprocessed def tokenize(self, texts: list[str] | list[dict] | list[tuple[str, str]], **kwargs) -> dict[str, Tensor]: """ .. deprecated:: `tokenize` is deprecated. Use `preprocess` instead. """ logger.warning_once("The `tokenize` method is deprecated, please use `preprocess` instead.") return self.preprocess(inputs=texts, **kwargs) def is_singular_input(self, inputs: Any) -> bool: """ Check if the input represents a single example or a batch of examples. Args: inputs: The input to check. Returns: bool: True if the input is a single example, False if it is a batch. """ list_types = (list, tuple) if is_datasets_available(): try: from datasets import Column list_types += (Column,) except ImportError: pass if isinstance(inputs, list_types): return False # Numpy arrays of unicode strings or objects are batches if isinstance(inputs, np.ndarray) and inputs.ndim >= 1 and inputs.dtype.kind in ("U", "O"): return False return True def save( self, path: str, model_name: str | None = None, create_model_card: bool = True, train_datasets: list[str] | None = None, safe_serialization: bool = True, ) -> None: """ Saves a model and its configuration files to a directory, so that it can be loaded again. Args: path (str): Path on disk where the model will be saved. model_name (str, optional): Optional model name. create_model_card (bool, optional): If True, create a README.md with basic information about this model. train_datasets (List[str], optional): Optional list with the names of the datasets used to train the model. safe_serialization (bool, optional): If True, save the model using safetensors. If False, save the model the traditional (but unsafe) PyTorch way. """ if path is None: return os.makedirs(path, exist_ok=True) logger.info(f"Saving model to {path}") modules_config = [] # Save model-level configuration options config = self._get_model_config() with open(os.path.join(path, "config_sentence_transformers.json"), "w", encoding="utf8") as fOut: json.dump(config, fOut, indent=2, sort_keys=True) # Save modules for idx, name in enumerate(self._modules): module: Module = self._modules[name] if ( idx == 0 and hasattr(module, "save_in_root") and module.save_in_root ): # Save first module in the main folder model_path = os.path.join(path, "") else: model_path = os.path.join(path, str(idx) + "_" + type(module).__name__) os.makedirs(model_path, exist_ok=True) # Try to save with safetensors, but fall back to the traditional PyTorch way if the module doesn't support it try: module.save(model_path, safe_serialization=safe_serialization) except TypeError: module.save(model_path) class_ref = type(module).__module__ # For remote modules, we want to remove "transformers_modules.{repo_name}": if class_ref.startswith("transformers_modules."): class_file = sys.modules[class_ref].__file__ # Save the custom module file dest_file = Path(model_path) / (Path(class_file).name) shutil.copy(class_file, dest_file) # Save all files imported in the custom module file for needed_file in get_relative_import_files(class_file): dest_file = Path(model_path) / (Path(needed_file).name) shutil.copy(needed_file, dest_file) # For remote modules, we want to ignore the "transformers_modules.{repo_id}" part, # i.e. we only want the filename class_ref = f"{class_ref.split('.')[-1]}.{type(module).__name__}" else: class_ref = f"{class_ref}.{type(module).__name__}" module_config = {"idx": idx, "name": name, "path": os.path.basename(model_path), "type": class_ref} if self.module_kwargs and name in self.module_kwargs and (module_kwargs := self.module_kwargs[name]): module_config["kwargs"] = module_kwargs modules_config.append(module_config) with open(os.path.join(path, "modules.json"), "w", encoding="utf8") as fOut: json.dump(modules_config, fOut, indent=2) if create_model_card: self._create_model_card(path, model_name, train_datasets) def _get_model_config(self) -> dict[str, Any]: return { "model_type": self.model_type, "__version__": { "sentence_transformers": __version__, "transformers": transformers.__version__, "pytorch": torch.__version__, }, "prompts": self.prompts, "default_prompt_name": self.default_prompt_name, } def save_pretrained( self, path: str, model_name: str | None = None, create_model_card: bool = True, train_datasets: list[str] | None = None, safe_serialization: bool = True, ) -> None: """ Saves a model and its configuration files to a directory, so that it can be loaded again. Args: path (str): Path on disk where the model will be saved. model_name (str, optional): Optional model name. create_model_card (bool, optional): If True, create a README.md with basic information about this model. train_datasets (List[str], optional): Optional list with the names of the datasets used to train the model. safe_serialization (bool, optional): If True, save the model using safetensors. If False, save the model the traditional (but unsafe) PyTorch way. """ self.save( path, model_name=model_name, create_model_card=create_model_card, train_datasets=train_datasets, safe_serialization=safe_serialization, ) def _update_default_model_id(self, model_card: str) -> str: """Update the default model ID in the model card.""" if self.model_card_data.model_id: model_card = model_card.replace( f'model = {self.__class__.__name__}("{self._model_card_model_id_placeholder}"', f'model = {self.__class__.__name__}("{self.model_card_data.model_id}"', ) return model_card def _create_model_card( self, path: str, model_name: str | None = None, train_datasets: list[str] | None = "deprecated" ) -> None: """ Create an automatic model card and store it in the specified path. Args: path (str): The path where the model card will be stored. model_name (Optional[str], optional): The name of the model. Defaults to None. train_datasets (Optional[List[str]], optional): Deprecated argument, ignored. Returns: None """ if model_name: model_path = Path(model_name) if not model_path.exists() and not self.model_card_data.model_id: self.model_card_data.model_id = model_name # Set the save directory so that assets (images, audio, etc.) can be saved alongside the model self.model_card_data.save_dir = path # If we loaded a model from the Hub, and no training was done, then # we don't generate a new model card, but reuse the old one instead. if self._model_card_text and "generated_from_trainer" not in self.model_card_data.tags: model_card = self._model_card_text model_card = self._update_default_model_id(model_card) else: try: model_card = generate_model_card(self) except Exception: logger.error( f"Error while generating model card:\n{traceback.format_exc()}" "Consider opening an issue on https://github.com/huggingface/sentence-transformers/issues with this traceback.\n" "Skipping model card creation." ) return with open(os.path.join(path, "README.md"), "w", encoding="utf8") as fOut: fOut.write(model_card) @save_to_hub_args_decorator def save_to_hub( self, repo_id: str, organization: str | None = None, token: str | None = None, private: bool | None = None, safe_serialization: bool = True, commit_message: str = "Add new model.", local_model_path: str | None = None, exist_ok: bool = False, replace_model_card: bool = False, train_datasets: list[str] | None = None, ) -> str: """ DEPRECATED, use `push_to_hub` instead. Uploads all elements of this model to a new HuggingFace Hub repository. Args: repo_id (str): Repository name for your model in the Hub, including the user or organization. token (str, optional): An authentication token (See https://huggingface.co/settings/token) private (bool, optional): Set to true, for hosting a private model safe_serialization (bool, optional): If true, save the model using safetensors. If false, save the model the traditional PyTorch way commit_message (str, optional): Message to commit while pushing. local_model_path (str, optional): Path of the model locally. If set, this file path will be uploaded. Otherwise, the current model will be uploaded exist_ok (bool, optional): If true, saving to an existing repository is OK. If false, saving only to a new repository is possible replace_model_card (bool, optional): If true, replace an existing model card in the hub with the automatically created model card train_datasets (List[str], optional): Datasets used to train the model. If set, the datasets will be added to the model card in the Hub. Returns: str: The url of the commit of your model in the repository on the Hugging Face Hub. """ logger.warning( "The `save_to_hub` method is deprecated and will be removed in a future version of SentenceTransformers." " Please use `push_to_hub` instead for future model uploads." ) if organization: if "/" not in repo_id: logger.warning( f'Providing an `organization` to `save_to_hub` is deprecated. Please use `repo_id="{organization}/{repo_id}"` instead.' ) repo_id = f"{organization}/{repo_id}" elif repo_id.split("/")[0] != organization: raise ValueError( "Providing an `organization` to `save_to_hub` is deprecated. Please use `repo_id` instead." ) else: logger.warning( f'Providing an `organization` to `save_to_hub` is deprecated. Please use `repo_id="{repo_id}"` instead.' ) return self.push_to_hub( repo_id=repo_id, token=token, private=private, safe_serialization=safe_serialization, commit_message=commit_message, local_model_path=local_model_path, exist_ok=exist_ok, replace_model_card=replace_model_card, train_datasets=train_datasets, ) def push_to_hub( self, repo_id: str, token: str | None = None, private: bool | None = None, safe_serialization: bool = True, commit_message: str | None = None, local_model_path: str | None = None, exist_ok: bool = False, replace_model_card: bool = False, train_datasets: list[str] | None = None, revision: str | None = None, create_pr: bool = False, ) -> str: """ Uploads all elements of this model to a HuggingFace Hub repository, creating it if it doesn't exist. Args: repo_id (str): Repository name for your model in the Hub, including the user or organization. token (str, optional): An authentication token (See https://huggingface.co/settings/token) private (bool, optional): Set to true, for hosting a private model safe_serialization (bool, optional): If true, save the model using safetensors. If false, save the model the traditional PyTorch way commit_message (str, optional): Message to commit while pushing. local_model_path (str, optional): Path of the model locally. If set, this file path will be uploaded. Otherwise, the current model will be uploaded exist_ok (bool, optional): If true, saving to an existing repository is OK. If false, saving only to a new repository is possible replace_model_card (bool, optional): If true, replace an existing model card in the hub with the automatically created model card. If false (default), keep the existing model card if one exists in the repository. train_datasets (List[str], optional): Datasets used to train the model. If set, the datasets will be added to the model card in the Hub. revision (str, optional): Branch to push the uploaded files to create_pr (bool, optional): If True, create a pull request instead of pushing directly to the main branch Returns: str: The url of the commit of your model in the repository on the Hugging Face Hub. """ api = HfApi(token=token) repo_url = api.create_repo( repo_id=repo_id, private=private, repo_type=None, exist_ok=exist_ok or create_pr, ) repo_id = repo_url.repo_id # Update the repo_id in case the old repo_id didn't contain a user or organization self.model_card_data.set_model_id(repo_id) if revision is not None: api.create_branch(repo_id=repo_id, branch=revision, exist_ok=True) if commit_message is None: backend = self.get_backend() if backend == "torch": commit_message = f"Add new {self.__class__.__name__} model" else: commit_message = f"Add new {self.__class__.__name__} model with an {backend} backend" commit_description = "" if create_pr: commit_description = f"""\ Hello! This pull request has been automatically generated to add {self.__class__.__name__} compatibility. ## Full Model Architecture: ``` {self} ``` {self._push_to_hub_usage_tip(repo_id)}""" if local_model_path: folder_url = api.upload_folder( repo_id=repo_id, folder_path=local_model_path, commit_message=commit_message, commit_description=commit_description if create_pr else None, create_pr=create_pr, revision=revision, ) else: with tempfile.TemporaryDirectory() as tmp_dir: if replace_model_card: create_model_card_for_path = True else: # If replace_model_card=False, skip model card creation only if there's already a README.md existing_readme = load_file_path( repo_id, "README.md", token=token, revision=revision, local_files_only=False ) create_model_card_for_path = existing_readme is None self.save( tmp_dir, model_name=repo_id, create_model_card=create_model_card_for_path, train_datasets=train_datasets, safe_serialization=safe_serialization, ) folder_url = api.upload_folder( repo_id=repo_id, folder_path=tmp_dir, commit_message=commit_message, commit_description=commit_description if create_pr else None, create_pr=create_pr, revision=revision, ) if create_pr: logger.info(f"A pull request has been created at {folder_url.pr_url}") return folder_url.pr_url return folder_url.commit_url def _load_modules( self, model_name_or_path: str, token: bool | str | None, cache_folder: str | None, revision: str | None = None, trust_remote_code: bool = False, local_files_only: bool = False, model_kwargs: dict[str, Any] | None = None, processor_kwargs: dict[str, Any] | None = None, config_kwargs: dict[str, Any] | None = None, ) -> tuple[list[nn.Module] | OrderedDict[str, nn.Module], dict[str, Any]]: load_kwargs = { "token": token, "cache_folder": cache_folder, "revision": revision, "trust_remote_code": trust_remote_code, "local_files_only": local_files_only, "model_kwargs": model_kwargs, "processor_kwargs": processor_kwargs, "config_kwargs": config_kwargs, } # Check if this is a Sentence Transformer model modules_json_path = load_file_path( model_name_or_path, "modules.json", token=token, cache_folder=cache_folder, revision=revision, local_files_only=local_files_only, ) if modules_json_path is None: logger.info(f"No modules.json found for {model_name_or_path}, initializing a new {self.model_type} model.") return self._load_default_modules(model_name_or_path, **load_kwargs) model_type_being_loaded = self._get_model_type( model_name_or_path, token=token, cache_folder=cache_folder, revision=revision, local_files_only=local_files_only, ) if model_type_being_loaded == self.model_type: logger.info(f"Loading {self.model_type} model from {model_name_or_path}.") return self._load_config_modules(model_name_or_path, **load_kwargs) logger.info(f"Converting {model_type_being_loaded} model {model_name_or_path} to {self.model_type}.") return self._load_converted_modules(model_name_or_path, **load_kwargs, model_type=model_type_being_loaded) @abstractmethod def _load_default_modules( self, model_name_or_path: str, token: bool | str | None, cache_folder: str | None, revision: str | None = None, trust_remote_code: bool = False, local_files_only: bool = False, model_kwargs: dict[str, Any] | None = None, processor_kwargs: dict[str, Any] | None = None, config_kwargs: dict[str, Any] | None = None, ) -> tuple[list[nn.Module] | OrderedDict[str, nn.Module], dict[str, Any]]: """ Args: model_name_or_path (str): The name or path of the pre-trained model. token (Optional[Union[bool, str]]): The token to use for the model. cache_folder (Optional[str]): The folder to cache the model. revision (Optional[str], optional): The revision of the model. Defaults to None. trust_remote_code (bool, optional): Whether to trust remote code. Defaults to False. local_files_only (bool, optional): Whether to use only local files. Defaults to False. model_kwargs (Optional[Dict[str, Any]], optional): Additional keyword arguments for the model. Defaults to None. processor_kwargs (Optional[Dict[str, Any]], optional): Additional keyword arguments for the processor/tokenizer. Defaults to None. config_kwargs (Optional[Dict[str, Any]], optional): Additional keyword arguments for the config. Defaults to None. Returns: List[nn.Module]: A list containing the transformer model and the pooling model. """ def _load_config_modules( self, model_name_or_path: str, token: bool | str | None, cache_folder: str | None, revision: str | None = None, trust_remote_code: bool = False, local_files_only: bool = False, model_kwargs: dict[str, Any] | None = None, processor_kwargs: dict[str, Any] | None = None, config_kwargs: dict[str, Any] | None = None, ) -> tuple[list[nn.Module] | OrderedDict[str, nn.Module], dict[str, Any]]: """ Loads a full model using the modules.json file. Args: model_name_or_path (str): The name or path of the pre-trained model. token (Optional[Union[bool, str]]): The token to use for the model. cache_folder (Optional[str]): The folder to cache the model. revision (Optional[str], optional): The revision of the model. Defaults to None. trust_remote_code (bool, optional): Whether to trust remote code. Defaults to False. local_files_only (bool, optional): Whether to use only local files. Defaults to False. model_kwargs (Optional[Dict[str, Any]], optional): Additional keyword arguments for the model. Defaults to None. processor_kwargs (Optional[Dict[str, Any]], optional): Additional keyword arguments for the processor/tokenizer. Defaults to None. config_kwargs (Optional[Dict[str, Any]], optional): Additional keyword arguments for the config. Defaults to None. Returns: Tuple[OrderedDict[str, nn.Module], OrderedDict[str, Any]]: An ordered dictionary containing the modules of the model and their kwargs. """ # Check if the config_sentence_transformers.json file exists (exists since v2 of the framework) config_sentence_transformers_json_path = load_file_path( model_name_or_path, "config_sentence_transformers.json", token=token, cache_folder=cache_folder, revision=revision, local_files_only=local_files_only, ) if config_sentence_transformers_json_path is not None: with open(config_sentence_transformers_json_path, encoding="utf8") as fIn: model_config = json.load(fIn) if ( "__version__" in model_config and "sentence_transformers" in model_config["__version__"] and version.parse(model_config["__version__"]["sentence_transformers"]) > version.parse(__version__) ): logger.warning( f"This model was created with Sentence Transformers version {model_config['__version__']['sentence_transformers']}, " f"but you're using version {__version__}. Consider updating to the latest version to avoid potential issues." ) self._parse_model_config(model_config) # Check if a readme exists model_card_path = load_file_path( model_name_or_path, "README.md", token=token, cache_folder=cache_folder, revision=revision, local_files_only=local_files_only, ) if model_card_path is not None: try: with open(model_card_path, encoding="utf8") as fIn: self._model_card_text = fIn.read() except Exception: pass # Load the modules modules_json_path = load_file_path( model_name_or_path, "modules.json", token=token, cache_folder=cache_folder, revision=revision, local_files_only=local_files_only, ) with open(modules_json_path, encoding="utf8") as fIn: modules_config = json.load(fIn) modules = OrderedDict() module_kwargs = OrderedDict() for module_config in modules_config: class_ref = module_config["type"] module_class: Module = self._load_module_class_from_ref( class_ref, model_name_or_path, trust_remote_code, revision, model_kwargs ) # Backwards compatibility: if the module is older and its `load` method only supports one parameter, # a path to a local directory containing the module files, then we load it with the old style load_signature = inspect.signature(module_class.load) # Check if the `load` method only accepts a single parameter (the path to the local directory). # This indicates an older module that does not support the newer loading method with multiple arguments. if len(load_signature.parameters) == 1: signature = inspect.signature(module_class.__init__) # Detect Transformer-based modules by checking for model/config kwargs in __init__. # Old custom modules (e.g. jinaai/jina-embeddings-v3) use model_args/config_args; # new-style modules use model_kwargs/config_kwargs. init_params = set(signature.parameters) uses_old_names = {"model_args", "config_args", "tokenizer_args"} & init_params uses_new_names = {"model_kwargs", "config_kwargs"} <= init_params if uses_new_names or uses_old_names: init_kwargs = Transformer._load_init_kwargs( model_name_or_path, # Loading-specific keyword arguments subfolder=module_config["path"], token=token, cache_folder=cache_folder, revision=revision, local_files_only=local_files_only, # Module-specific keyword arguments trust_remote_code=trust_remote_code, model_kwargs=model_kwargs, processor_kwargs=processor_kwargs, config_kwargs=config_kwargs, backend=self.backend, ) # Remap new-style keys back to old-style for old custom modules. new_to_old_name_mapping = { "model_kwargs": "model_args", "processor_kwargs": "tokenizer_args", "config_kwargs": "config_args", } if uses_old_names and not uses_new_names: for new_name, old_name in new_to_old_name_mapping.items(): if new_name in init_kwargs: init_kwargs[old_name] = init_kwargs.pop(new_name) # Some new config value (defaults) might not be supported by old modules, so we can drop them optional_params = {"modality_config", "module_output_name"} for new_name in optional_params: if new_name in init_kwargs and new_name not in init_params: init_kwargs.pop(new_name) module = module_class(model_name_or_path, **init_kwargs) else: # Old modules that don't support the new loading method and don't seem Transformer-based # are loaded by downloading the full directories and calling .load() with the old style # (i.e. only a path to the local directory) local_path = load_dir_path( model_name_or_path=model_name_or_path, subfolder=module_config["path"], token=token, cache_folder=cache_folder, revision=revision, local_files_only=local_files_only, ) module = module_class.load(local_path) else: # Newer modules that support the new loading method are loaded with the new style # i.e. with many keyword arguments that can optionally be used by the modules module = module_class.load( model_name_or_path, # Loading-specific keyword arguments subfolder=module_config["path"], token=token, cache_folder=cache_folder, revision=revision, local_files_only=local_files_only, # Module-specific keyword arguments trust_remote_code=trust_remote_code, model_kwargs=model_kwargs, processor_kwargs=processor_kwargs, config_kwargs=config_kwargs, backend=self.backend, ) modules[module_config["name"]] = module module_kwargs[module_config["name"]] = module_config.get("kwargs", []) if revision is None: path_parts = Path(modules_json_path) if len(path_parts.parts) >= 2: revision_path_part = Path(modules_json_path).parts[-2] if len(revision_path_part) == 40 and all(c in "0123456789abcdef" for c in revision_path_part): revision = revision_path_part if not local_files_only: self.model_card_data.set_base_model(model_name_or_path, revision=revision) return modules, module_kwargs def _parse_model_config(self, model_config: dict[str, Any]) -> None: """Parse model configuration and merge saved prompts/defaults with user-provided values. User-provided prompts and default_prompt_name take precedence over saved config values. Saved prompts are only used for keys not already present in ``self.prompts``, or where the current value is ``None`` (i.e. a default placeholder, not yet filled by the user or config). Empty string ``""`` is treated as an intentional user-provided value and will not be overwritten. """ # Only update prompts that aren't already set by the user or defaults for prompt_name, prompt_text in model_config.get("prompts", {}).items(): if prompt_name not in self.prompts or self.prompts[prompt_name] is None: self.prompts[prompt_name] = prompt_text if self.default_prompt_name is None: self.default_prompt_name = model_config.get("default_prompt_name", None) def _load_converted_modules( self, model_name_or_path: str, token: bool | str | None, cache_folder: str | None, revision: str | None = None, trust_remote_code: bool = False, local_files_only: bool = False, model_kwargs: dict[str, Any] | None = None, processor_kwargs: dict[str, Any] | None = None, config_kwargs: dict[str, Any] | None = None, model_type: str | None = None, ) -> tuple[list[nn.Module] | OrderedDict[str, nn.Module], dict[str, Any]]: return self._load_default_modules( model_name_or_path, token=token, cache_folder=cache_folder, revision=revision, trust_remote_code=trust_remote_code, local_files_only=local_files_only, model_kwargs=model_kwargs, processor_kwargs=processor_kwargs, config_kwargs=config_kwargs, ) def _get_model_type( self, model_name_or_path: str, token: bool | str | None, cache_folder: str | None, revision: str | None = None, local_files_only: bool = False, ) -> str: """ Retrieves the model_type from the config_sentence_transformers.json file. This is used to determine whether the model being loaded matches the current class (e.g., a SentenceTransformer model loaded with SentenceTransformer, or a SparseEncoder model loaded with SparseEncoder). When the model type doesn't match, we switch to a converted loading method to ensure compatibility. Defaults to "SentenceTransformer" if the config file is missing or has no "model_type" key, for backwards compatibility with older models. Args: model_name_or_path (str): The name or path of the pre-trained model. token (Optional[Union[bool, str]]): The token to use for the model. cache_folder (Optional[str]): The folder to cache the model. revision (Optional[str], optional): The revision of the model. Defaults to None. local_files_only (bool, optional): Whether to use only local files. Defaults to False. Returns: str: The model type, e.g. "SentenceTransformer", "SparseEncoder", or "CrossEncoder". """ config_sentence_transformers_json_path = load_file_path( model_name_or_path, "config_sentence_transformers.json", token=token, cache_folder=cache_folder, revision=revision, local_files_only=local_files_only, ) if config_sentence_transformers_json_path is None: return "SentenceTransformer" with open(config_sentence_transformers_json_path, encoding="utf8") as fIn: config = json.load(fIn) # Older SentenceTransformer models won't have "model_type", so those default to "SentenceTransformer" return config.get("model_type", "SentenceTransformer") def _load_module_class_from_ref( self, class_ref: str, model_name_or_path: str, trust_remote_code: bool, revision: str | None, model_kwargs: dict[str, Any] | None, ) -> nn.Module: """ Load a module class from a class reference string. Args: class_ref: The class reference string (e.g., "sentence_transformers.sentence_transformer.modules.Pooling") model_name_or_path: The model name or path trust_remote_code: Whether to trust remote code revision: The model revision model_kwargs: Additional model kwargs Returns: The module class """ # If the class is from sentence_transformers, we can directly import it, # otherwise, we try to import it dynamically, and if that fails, we fall back to the default import if class_ref.startswith("sentence_transformers."): return import_from_string(class_ref) if trust_remote_code or os.path.exists(model_name_or_path): code_revision = model_kwargs.pop("code_revision", None) if model_kwargs else None try: return get_class_from_dynamic_module( class_ref, model_name_or_path, revision=revision, code_revision=code_revision, ) except (OSError, ValueError): # Ignore the error if 1) the file does not exist, or 2) the class_ref is not correctly formatted/found pass return import_from_string(class_ref) def evaluate(self, evaluator: BaseEvaluator, output_path: str | None = None) -> dict[str, float] | float: """ Evaluate the model based on an evaluator Args: evaluator (BaseEvaluator): The evaluator used to evaluate the model. output_path (str, optional): The path where the evaluator can write the results. Defaults to None. Returns: The evaluation results. """ if output_path is not None: os.makedirs(output_path, exist_ok=True) return evaluator(self, output_path) def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs: dict[str, Any] | None = None) -> None: """Enable gradient checkpointing for the model.""" # Propagate the gradient checkpointing to the transformer model for module in self.modules(): if module is not self and hasattr(module, "gradient_checkpointing_enable"): try: module.gradient_checkpointing_enable(gradient_checkpointing_kwargs) except TypeError: module.gradient_checkpointing_enable() @property def device(self) -> torch.device: """ Get torch.device from module, assuming that the whole module has one device. In case there are no PyTorch parameters, fall back to CPU. """ if (transformers_model := self.transformers_model) is not None and hasattr(transformers_model, "device"): return transformers_model.device if len(self._modules) and hasattr(self[0], "auto_model") and hasattr(self[0].auto_model, "device"): return self[0].auto_model.device try: return next(self.parameters()).device except StopIteration: # Fallback for nn.DataParallel compatibility when parameters() is empty def find_tensor_attributes(module: nn.Module) -> list[tuple[str, Tensor]]: tuples = [(k, v) for k, v in module.__dict__.items() if torch.is_tensor(v)] return tuples gen = self._named_members(get_members_fn=find_tensor_attributes) try: first_tuple = next(gen) return first_tuple[1].device except StopIteration: return torch.device("cpu") def start_multi_process_pool( self, target_devices: list[str] | None = None ) -> dict[Literal["input", "output", "processes"], Any]: """ Starts a multi-process pool to infer with several independent processes. This method is recommended if you want to predict on multiple GPUs or CPUs. It is advised to start only one process per GPU. This method works together with predict and stop_multi_process_pool. Args: target_devices (List[str], optional): PyTorch target devices, e.g. ["cuda:0", "cuda:1", ...], ["npu:0", "npu:1", ...], or ["cpu", "cpu", "cpu", "cpu"]. If target_devices is None and CUDA/NPU is available, then all available CUDA/NPU devices will be used. If target_devices is None and CUDA/NPU is not available, then 4 CPU devices will be used. Returns: Dict[str, Any]: A dictionary with the target processes, an input queue, and an output queue. """ if target_devices is None: if torch.cuda.is_available(): target_devices = [f"cuda:{i}" for i in range(torch.cuda.device_count())] elif is_torch_npu_available(): target_devices = [f"npu:{i}" for i in range(torch.npu.device_count())] else: logger.info("CUDA/NPU is not available. Starting 4 CPU workers") target_devices = ["cpu"] * 4 logger.info(f"Starting multi-process pool on devices: {', '.join(map(str, target_devices))}") # Move model to CPU and share memory so child processes can access it. # Note: this modifies the model in-place, the model will remain on CPU after this call. self.to("cpu") self.share_memory() ctx = mp.get_context("spawn") input_queue = ctx.Queue() output_queue = ctx.Queue() processes = [] for device_id in target_devices: p = ctx.Process( target=self.__class__._multi_process_worker, args=(device_id, self, input_queue, output_queue), daemon=True, ) p.start() processes.append(p) return {"input": input_queue, "output": output_queue, "processes": processes} @staticmethod def stop_multi_process_pool(pool: dict[Literal["input", "output", "processes"], Any]) -> None: """ Stops all processes started with start_multi_process_pool. Args: pool (Dict[str, object]): A dictionary containing the input queue, output queue, and process list. Returns: None """ for p in pool["processes"]: p.terminate() for p in pool["processes"]: p.join() p.close() pool["input"].close() pool["output"].close() def _multi_process(self, *args, **kwargs): raise NotImplementedError("This method should be implemented in subclasses.") @staticmethod def _multi_process_worker( target_device: str, model: BaseModel, input_queue: Queue, results_queue: Queue, ) -> None: """Worker function for multi-process inference. Must be overridden by subclasses. This is called as the target function in each spawned process by :meth:`start_multi_process_pool`. Subclasses should implement this to read from ``input_queue``, run inference on ``target_device``, and write results to ``results_queue``. """ raise NotImplementedError("This method should be implemented in subclasses.") @property def tokenizer(self) -> Any: """ Property to get the tokenizer that is used by this model """ return self[0].tokenizer @tokenizer.setter def tokenizer(self, value) -> None: """ Property to set the tokenizer that should be used by this model """ try: self[0].tokenizer = value except AttributeError: raise AttributeError( f"The first module ({type(self[0]).__name__}) does not have a 'tokenizer' attribute." ) from None @property def processor(self) -> Any: """ Property to get the processor that is used by this model """ return self[0].processor @property def max_seq_length(self) -> int | None: """ Returns the maximal input sequence length for the model. Longer inputs will be truncated. Returns: Optional[int]: The maximal input sequence length, or None if not defined. """ return getattr(self[0], "max_seq_length", None) @max_seq_length.setter def max_seq_length(self, value) -> None: """ Property to set the maximal input sequence length for the model. Longer inputs will be truncated. """ self[0].max_seq_length = value @property def transformers_model(self) -> PreTrainedModel | None: """ Property to get the underlying transformers PreTrainedModel instance, if it exists. Note that it's possible for a model to have multiple underlying transformers models, but this property will return the first one it finds in the module hierarchy. .. note:: This property can also return e.g. ORTModelForFeatureExtraction or OVModelForFeatureExtraction instances from the optimum-intel and optimum-onnx libraries, if the model is loaded using ``backend="onnx"`` or ``backend="openvino"``. Returns: PreTrainedModel or None: The underlying transformers model or None if not found. """ for module in self.modules(): # The Transformer check allows for returning underlying models with backend="onnx" or "openvino" if isinstance(module, Transformer): return module.model if isinstance(module, PreTrainedModel): return module return None @property def _target_device(self) -> torch.device: logger.warning( f"`{self.__class__.__name__}._target_device` has been deprecated. Please use `{self.__class__.__name__}.device` instead.", ) return self.device @_target_device.setter def _target_device(self, device: int | str | torch.device | None = None) -> None: logger.warning( f"`{self.__class__.__name__}._target_device` has been deprecated. Please use `to(device)` instead.", ) self.to(device) @property def dtype(self) -> torch.dtype | None: """ `torch.dtype`: The dtype of the module (assuming that all the module parameters have the same dtype). """ try: return next(self.parameters()).dtype except StopIteration: return None @property def _no_split_modules(self) -> list[str]: """ Return the list of modules that should not be split when using model parallelism. """ return [] @property def _keys_to_ignore_on_save(self) -> list[str]: """ Return the list of keys to ignore when saving the model. """ return [] def _push_to_hub_usage_tip(self, repo_id: str) -> str: """Return a usage tip snippet for the push_to_hub PR description. Subclasses can override this to provide model-type-specific example code. """ class_name = self.__class__.__name__ backend = self.get_backend() return f"""\ ## Testing this pull request You can test this pull request before merging by loading the model from this PR with the `revision` argument: ```python from sentence_transformers import {class_name} # NOTE: Update this to the number of your pull request pr_number = 2 model = {class_name}( "{repo_id}", revision=f"refs/pr/{{pr_number}}", backend="{backend}", ) ``` """