from __future__ import annotations import json import os from pathlib import Path try: from typing import Self except ImportError: from typing_extensions import Self from torch import Tensor, nn from transformers.utils import logging from sentence_transformers.base.modality import format_modality, infer_batch_modality from sentence_transformers.base.modality_types import MODALITY_TO_PROCESSOR_ARG, Modality, PairInput, SingleInput from sentence_transformers.base.modules.input_module import InputModule from sentence_transformers.base.modules.module import Module from sentence_transformers.util import import_from_string, load_dir_path logger = logging.get_logger(__name__) class Router(InputModule): forward_kwargs = {"task", "modality"} config_keys: list[str] = ["default_route", "allow_empty_key", "route_mappings"] config_file_name = "router_config.json" def __init__( self, sub_modules: dict[str, list[Module]], default_route: str | None = None, allow_empty_key: bool = True, route_mappings: dict[tuple[str | None, str | tuple[str, ...] | None], str] | None = None, ) -> None: r""" This model allows creating flexible SentenceTransformer models that dynamically route inputs to different processing modules based on: 1. Task type (e.g., "query" or "document") for asymmetric retrieval models 2. Modality (e.g., "text", "image", or ("text", "image")) for crossmodal or multimodal models 3. Combination of both for complex routing scenarios Tips: - The ``task`` argument in ``model.encode()`` specifies which route to use - ``model.encode_query()`` and ``model.encode_document()`` are convenient shorthands for ``task="query"`` and ``task="document"`` - Modality is automatically inferred from input data (text strings, PIL Images, etc.) - You can override automatic inference by passing ``modality`` in ``model.encode()`` (and its variants) explicitly Route Priority: 1. Exact match: ``(task, modality)`` - e.g., ``("query", "text")`` 2. Task with any modality: ``(task, None)`` - e.g., ``("query", None)`` 3. Any task with modality: ``(None, modality)`` - e.g., ``(None, "image")`` 4. Catch-all: ``(None, None)`` 5. Direct lookup by task name in ``sub_modules`` 6. Direct lookup by modality name in ``sub_modules`` 7. Fall back to ``default_route`` if set In the below examples, the ``Router`` model is used to create asymmetric models with different encoders for queries and documents. In these examples, the "query" route is efficient (e.g., using SparseStaticEmbedding), while the "document" route uses a more complex model (e.g. a Transformers module). This allows for efficient query encoding while still using a powerful document encoder, but the combinations are not limited to this. Example: :: from sentence_transformers import SentenceTransformer from sentence_transformers.sentence_transformer.modules import Router, Normalize # Use a regular SentenceTransformer for the document embeddings, and a static embedding model for the query embeddings document_embedder = SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1") query_embedder = SentenceTransformer("sentence-transformers/static-retrieval-mrl-en-v1") router = Router.for_query_document( query_modules=list(query_embedder.children()), document_modules=list(document_embedder.children()), ) normalize = Normalize() # Create an asymmetric model with different encoders for queries and documents model = SentenceTransformer( modules=[router, normalize], ) # ... requires more training to align the vector spaces # Use the query & document routes query_embedding = model.encode_query("What is the capital of France?") document_embedding = model.encode_document("Paris is the capital of France.") :: from sentence_transformers.sparse_encoder.modules import Router, SparseStaticEmbedding, SpladePooling, Transformer from sentence_transformers.sparse_encoder import SparseEncoder # Load an asymmetric model with different encoders for queries and documents doc_encoder = Transformer("opensearch-project/opensearch-neural-sparse-encoding-doc-v3-distill", transformer_task="fill-mask") router = Router.for_query_document( query_modules=[ SparseStaticEmbedding.from_json( "opensearch-project/opensearch-neural-sparse-encoding-doc-v3-distill", tokenizer=doc_encoder.tokenizer, frozen=True, ), ], document_modules=[ doc_encoder, SpladePooling(pooling_strategy="max", activation_function="log1p_relu"), ], ) model = SparseEncoder(modules=[router], similarity_fn_name="dot") query = "What's the weather in ny now?" document = "Currently New York is rainy." query_embed = model.encode_query(query) document_embed = model.encode_document(document) sim = model.similarity(query_embed, document_embed) print(f"Similarity: {sim}") # Visualize top tokens for each text top_k = 10 print(f"Top tokens {top_k} for each text:") decoded_query = model.decode(query_embed, top_k=top_k) decoded_document = model.decode(document_embed) for i in range(min(top_k, len(decoded_query))): query_token, query_score = decoded_query[i] doc_score = next((score for token, score in decoded_document if token == query_token), 0) if doc_score != 0: print(f"Token: {query_token}, Query score: {query_score:.4f}, Document score: {doc_score:.4f}") ''' Similarity: tensor([[11.1105]], device='cuda:0') Top tokens 10 for each text: Token: ny, Query score: 5.7729, Document score: 0.8049 Token: weather, Query score: 4.5684, Document score: 0.9710 Token: now, Query score: 3.5895, Document score: 0.4720 Token: ?, Query score: 3.3313, Document score: 0.0286 Token: what, Query score: 2.7699, Document score: 0.0787 Token: in, Query score: 0.4989, Document score: 0.0417 ''' Multimodal Example: :: from PIL import Image from sentence_transformers import SentenceTransformer from sentence_transformers.sentence_transformer.modules import Dense, Pooling, Router, Transformer # Create separate encoders for different modalities text_encoder = Transformer("sentence-transformers/all-MiniLM-L6-v2") # Project to 768 dims to match image encoder text_dense = Dense(text_encoder.get_embedding_dimension(), 768, module_input_name="token_embeddings") image_encoder = Transformer( "ModernVBERT/modernvbert", model_kwargs={"trust_remote_code": True}, processor_kwargs={"trust_remote_code": True}, config_kwargs={"trust_remote_code": True}, ) pooling = Pooling(text_encoder.get_embedding_dimension()) # Route based on modality router = Router( sub_modules={ "text": [text_encoder, text_dense], "image": [image_encoder], }, route_mappings={ (None, "text"): "text", # Any task with text goes to text encoder (None, ("text", "image")): "image", # Any task with text-image together goes to image encoder }, ) model = SentenceTransformer(modules=[router, pooling]) # Modality is automatically inferred text_embedding = model.encode("A photo of a cat") multimodal_embedding = model.encode({"text": "A photo of a ", "image": Image.open("cat.jpg")}) # Compute the similarity; it'll be poor as the model hasn't yet been trained similarity = model.similarity(text_embedding, multimodal_embedding) Hybrid Asymmetric + Multimodal Example: :: from sentence_transformers import SentenceTransformer from sentence_transformers.sentence_transformer.modules import Router # Different encoders for query text, document text, and images router = Router( sub_modules={ "query_text": [query_text_modules], "doc_text": [document_text_modules], "image": [image_modules], }, route_mappings={ ("query", "text"): "query_text", # Query text uses efficient encoder ("document", "text"): "doc_text", # Document text uses powerful encoder (None, ("text", "image")): "image", # Any text-image together goes to image encoder }, ) model = SentenceTransformer(modules=[router]) # Explicit task + automatic modality inference query_embedding = model.encode_query("Find images of cats") doc_embedding = model.encode_document("Article about cats") multimodal_embedding = model.encode({"text": "A photo of a cat", "image": Image.open("cat.jpg")}) .. note:: When training models with the :class:`~sentence_transformers.base.modules.Router` module, you must use the ``router_mapping`` argument in the :class:`~sentence_transformers.sentence_transformer.training_args.SentenceTransformerTrainingArguments` or :class:`~sentence_transformers.sparse_encoder.training_args.SparseEncoderTrainingArguments` to map the training dataset columns to the correct route ("query" or "document"). For example, if your training dataset(s) have ``["question", "positive", "negative"]`` columns, then you can use the following mapping:: args = SparseEncoderTrainingArguments( ..., router_mapping={ "question": "query", "positive": "document", "negative": "document", } ) Additionally, it is common to use a different learning rate for the different routes. For this, you should use the ``learning_rate_mapping`` argument in the :class:`~sentence_transformers.sentence_transformer.training_args.SentenceTransformerTrainingArguments` or :class:`~sentence_transformers.sparse_encoder.training_args.SparseEncoderTrainingArguments` to map parameter patterns to their learning rates. For example, if you want to use a learning rate of ``1e-3`` for an SparseStaticEmbedding module and ``2e-5`` for the rest of the model, you can do this:: args = SparseEncoderTrainingArguments( ..., learning_rate=2e-5, learning_rate_mapping={ r"SparseStaticEmbedding\.*": 1e-3, } ) Args: sub_modules: Mapping of route keys to lists of modules. Each key corresponds to a specific route name (e.g., "text_query", "text_document", "image", "multimodal"). Each route contains a list of modules that will be applied sequentially when that route is selected. default_route: The default route to use if no task type or modality is specified. If None, an exception will be thrown if no task type is specified. If ``allow_empty_key`` is True, the first key in sub_modules will be used as the default route. Defaults to None. allow_empty_key: If True, allows the default route to be set to the first key in `sub_modules` if ``default_route`` is None. Defaults to True. route_mappings: Optional dictionary mapping (task, modality) tuples to route keys in sub_modules. This enables sophisticated routing logic based on combinations of task and modality: - Use ``None`` as a wildcard for either task or modality to create catch-all rules - Modality can be a string (e.g., ``"text"``, ``"image"``) or tuple (e.g., ``("text", "image")``) - Routes are resolved with a priority order (see **Route Resolution Priority** above) - All mapped routes must exist in ``sub_modules`` (validated at initialization) Example mappings:: { # Exact matches (highest priority) ("query", "text"): "efficient_text_encoder", ("document", "text"): "powerful_text_encoder", # Task with any modality ("query", None): "query_encoder", # All query tasks # Any task with specific modality (None, "image"): "image_encoder", # All image inputs (None, ("text", "image")): "multimodal_encoder", # Multimodal inputs # Catch-all (lowest priority) (None, None): "default_encoder", } If not provided, the router will attempt direct lookup using the task or modality as the route key in ``sub_modules``, then fall back to ``default_route``. """ super().__init__() if sub_modules is None or len(sub_modules) == 0: raise ValueError("The sub_modules dictionary cannot be empty.") if default_route is not None and default_route not in sub_modules: raise ValueError(f"Default route '{default_route}' not found in route keys: {list(sub_modules.keys())}") self.sub_modules = nn.ModuleDict( {route_name: nn.Sequential(*modules) for route_name, modules in sub_modules.items()} ) # Validate that all route_mappings point to existing sub_modules if route_mappings: for (task, modality), target_route in route_mappings.items(): if target_route not in sub_modules: raise ValueError( f"route_mappings contains mapping to '{target_route}' which is not in sub_modules. " f"Available routes: {list(sub_modules.keys())}" ) # If allow_empty_key is True, we can set a default route to the first key in sub_modules. if allow_empty_key and default_route is None: default_route = next(iter(sub_modules.keys())) self.default_route = default_route self.allow_empty_key = allow_empty_key # Normalize tuple modalities in route_mappings keys to sorted tuples so they match # the sorted tuples returned by infer_modality for multimodal dict inputs. self.route_mappings = ( { (task, tuple(sorted(modality)) if isinstance(modality, tuple) else modality): route for (task, modality), route in route_mappings.items() } if route_mappings else {} ) @property def modalities(self) -> list[Modality]: """The union of modalities supported by all sub-module input modules.""" return sorted( { modality for route in self.sub_modules.values() for modality in getattr(route[0], "modalities", ["text"]) }, key=str, ) def _get_routes_string(self): # Build concise but helpful error message routes = [f"task={name!r}" for name in self.sub_modules.keys()] for task, modality in self.route_mappings.keys(): if modality is None: routes.append(f"task={task!r}") elif task is None: routes.append(f"modality={format_modality(modality)!r}") else: routes.append(f"(task={task!r}, modality={format_modality(modality)!r})") if not routes: return "" elif len(routes) == 1: return routes[0] return ", ".join(routes[:-1]) + " and " + routes[-1] def _resolve_route_name( self, task: str | None = None, modality: str | tuple[str, ...] | None = None ) -> str | None: """ Resolve the route key based on task and modality. Args: task: The task type (e.g., "query", "document") modality: The modality (e.g., "text", "image", ("text", "image")) Returns: The resolved route key, or None if no route is found """ # Normalize tuple modality to sorted tuple to match the normalization applied at init time if isinstance(modality, tuple): modality = tuple(sorted(modality)) # Try exact match first: (task, modality) if (task, modality) in self.route_mappings: return self.route_mappings[(task, modality)] # Try task with any modality: (task, None) if (task, None) in self.route_mappings: return self.route_mappings[(task, None)] # Try any task with modality: (None, modality) if (None, modality) in self.route_mappings: return self.route_mappings[(None, modality)] # Try any task with any modality: (None, None) if (None, None) in self.route_mappings: return self.route_mappings[(None, None)] # Fallback: Try direct lookup by task if task and task in self.sub_modules: return task # Fallback: Try direct lookup by modality if modality and modality in self.sub_modules: return modality if task is not None: raise ValueError(f"No route found for task type '{task}'. Available routes: {self._get_routes_string()}") # No route found return None def _resolve_route(self, task: str | None = None, modality: str | tuple[str, ...] | None = None) -> str | None: route = self._resolve_route_name(task=task, modality=modality) # Fall back to default route if no route was found if route is None: route = self.default_route # If still no route, raise an error if route is None: if self.training: raise ValueError( "You must provide a `router_mapping` argument on the training arguments, " "or set a default route in the `Router` module." ) modality_display = format_modality(modality) if modality is not None else None error_msg = f"Could not determine route for task={task!r}, modality={modality_display!r}. " error_msg += f"Available routes: {self._get_routes_string()}. " error_msg += "Consider specifying the `task` parameter in `model.encode`, or setting a default route in the `Router`." raise ValueError(error_msg) if route not in self.sub_modules: raise ValueError( f"Resolved route '{route}' not found in sub_modules. Available submodule keys: {list(self.sub_modules.keys())}" ) return route @classmethod def for_query_document( cls, query_modules: list[Module], document_modules: list[Module], default_route: str | None = "document", allow_empty_key: bool = True, ) -> Self: """ Creates a Router model specifically for query and document modules, allowing convenient usage via `model.encode_query` and `model.encode_document`. Args: query_modules: List of modules to be applied for the "query" task type. document_modules: List of modules to be applied for the "document" task type. default_route: The default route to use if no task type is specified. If None, an exception will be thrown if no task type is specified. If ``allow_empty_key`` is True, the first key in sub_modules will be used as the default route. Defaults to "document". allow_empty_key: If True, allows the default route to be set to the first key in `sub_modules` if ``default_route`` is None. Defaults to True. Returns: Router: An instance of the Router model with the specified query and document modules. """ return cls( sub_modules={"query": query_modules, "document": document_modules}, default_route=default_route, allow_empty_key=allow_empty_key, ) def forward( self, features: dict[str, Tensor], task: str | None = None, modality: str | tuple[str, ...] | None = None, **kwargs, ) -> dict[str, Tensor]: """Route ``features`` through the resolved sub-module pipeline. Resolves the route from ``task`` and ``modality`` (falling back to values stored in ``features`` if not provided), then sequentially applies all modules in the matched route. Args: features: Input features dictionary (e.g. from :meth:`preprocess`). task: Task type used for routing (e.g. ``"query"``, ``"document"``). Falls back to ``features["task"]`` if not provided. modality: Modality used for routing (e.g. ``"text"``, ``"image"``). Falls back to ``features["modality"]`` if not provided. **kwargs: Extra keyword arguments forwarded to each sub-module's ``forward`` (filtered by each module's ``forward_kwargs``). Returns: The features dictionary after passing through all modules in the resolved route. """ # Get task from features if not provided if task is None: task = features.get("task", None) # Get modality from features if not provided if modality is None: modality = features.get("modality", None) # Resolve the route using task and modality route = self._resolve_route(task=task, modality=modality) # Pass task and modality to downstream modules kwargs["task"] = task kwargs["modality"] = modality for module in self.sub_modules[route]: module_kwargs = { key: value for key, value in kwargs.items() if hasattr(module, "forward_kwargs") and key in module.forward_kwargs } features = module(features, **module_kwargs) return features def __repr__(self) -> str: # Use nn.Module's repr to show the full sub-module structure, # rather than Module's repr which only shows config_keys return nn.Module.__repr__(self) def extra_repr(self) -> str: parts = [] if self.default_route is not None: parts.append(f"default_route={self.default_route!r}") if self.route_mappings: parts.append(f"route_mappings={self.route_mappings}") return ", ".join(parts) def get_embedding_dimension(self) -> int | None: dims = [] for sub_modules in self.sub_modules.values(): for module in reversed(sub_modules): # Fallback names for third-party modules that only define a deprecated name for name in ( "get_embedding_dimension", "get_sentence_embedding_dimension", "get_word_embedding_dimension", ): if hasattr(module, name): dims.append(getattr(module, name)()) break else: continue break if not dims: return None if len(set(dims)) > 1: logger.warning_once( f"Different embedding dimensions detected across routes: {set(dims)}. Using the first value." ) return dims[0] def save(self, output_path: str, safe_serialization: bool = True, **kwargs): model_lookup = {} model_types = {} model_structure = {} for name, models in self.sub_modules.items(): model_structure[name] = [] for module_idx, model in enumerate(models): model_id = f"{name}_{module_idx}_{type(model).__name__}" model_lookup[model_id] = model model_types[model_id] = f"{type(model).__module__}.{type(model).__name__}" model_structure[name].append(model_id) for model_id, model in model_lookup.items(): model_path = os.path.join(output_path, str(model_id)) os.makedirs(model_path, exist_ok=True) try: model.save(model_path, safe_serialization=safe_serialization, **kwargs) except TypeError: # Fallback for legacy models that do not support kwargs model.save(model_path) # Get config dict and convert route_mappings keys to strings for JSON serialization config_dict = self.get_config_dict() if "route_mappings" in config_dict and config_dict["route_mappings"]: # Convert tuple keys to string representation for JSON config_dict["route_mappings"] = {str(key): value for key, value in config_dict["route_mappings"].items()} with open(os.path.join(output_path, self.config_file_name), "w", encoding="utf8") as fOut: json.dump( { "types": model_types, "structure": model_structure, "parameters": config_dict, }, fOut, indent=4, ) def preprocess( self, inputs: list[SingleInput | PairInput], prompt: str | None = None, task: str | None = None, modality: Modality | None = None, **kwargs, ): """Resolve the route from ``task`` and ``modality``, then delegate preprocessing to the first module of the matched route. The returned dictionary includes ``"task"`` and (if available) ``"modality"`` keys so that :meth:`forward` can route without re-inference. """ # Backwards compatibility branch: for when texts are list of dicts with task types as keys if inputs and isinstance(inputs[0], dict) and task is None: # Extract the task type key from the dictionaries dict_keys = set(key for text in inputs for key in text.keys()) # If the input keys are not modalities, they might be task types from backwards compatibility if not (dict_keys & set(MODALITY_TO_PROCESSOR_ARG.keys())): if len(dict_keys) > 1: raise ValueError( "You cannot pass a list of dictionaries with different task types. " "Please ensure all dictionaries have the same task type key, or pass non-dictionary inputs " "while providing the `task` argument." ) # If there's just one task type, extract it and remove the dictionary structure task = dict_keys.pop() inputs = [sample[task] for sample in inputs] # Infer modality if not provided if modality is None and inputs: try: modality = infer_batch_modality(inputs, supported_modalities=self.modalities) except (ValueError, TypeError): # If modality inference fails, proceed without it pass # Resolve route route = self._resolve_route(task=task, modality=modality) input_module = self.sub_modules[route][0] tokenized = input_module.preprocess(inputs, prompt=prompt, **kwargs) tokenized["task"] = task if modality is not None: tokenized["modality"] = modality return tokenized @classmethod def load( cls, model_name_or_path: str, subfolder: str = "", token: bool | str | None = None, cache_folder: str | None = None, revision: str | None = None, local_files_only: bool = False, **kwargs, ) -> Self: hub_kwargs = { "token": token, "cache_folder": cache_folder, "revision": revision, "local_files_only": local_files_only, } # Try the official config file first, then fall back to the legacy config file config = cls.load_config(model_name_or_path=model_name_or_path, subfolder=subfolder, **hub_kwargs) if not config: config = cls.load_config( model_name_or_path=model_name_or_path, config_filename="config.json", subfolder=subfolder, **hub_kwargs ) modules = {} for model_id, model_type in config["types"].items(): module_class: Module = import_from_string(model_type) try: module = module_class.load( model_name_or_path, subfolder=Path(subfolder, model_id).as_posix(), **hub_kwargs, **kwargs ) except TypeError: local_path = load_dir_path( model_name_or_path=model_name_or_path, subfolder=Path(subfolder, model_id).as_posix(), **hub_kwargs ) module = module_class.load(local_path) modules[model_id] = module model_structure = {} for key_name, models_list in config["structure"].items(): model_structure[key_name] = [] for model_id in models_list: model_structure[key_name].append(modules[model_id]) # Convert route_mappings string keys back to tuples parameters = config["parameters"].copy() if "route_mappings" in parameters and parameters["route_mappings"]: route_mappings = {} for key_str, value in parameters["route_mappings"].items(): # Parse the string representation back to tuple # Format: "('task', 'modality')" or "(None, 'modality')" etc. try: import ast key_tuple = ast.literal_eval(key_str) route_mappings[key_tuple] = value except (ValueError, SyntaxError): logger.warning(f"Could not parse route_mapping key: {key_str}. Skipping.") parameters["route_mappings"] = route_mappings model = cls(model_structure, **parameters) return model @property def tokenizer(self): # We might have multiple tokenizers, one for each route, but we can only return one here. for sub_modules in self.sub_modules.values(): input_module: InputModule = sub_modules[0] if hasattr(input_module, "tokenizer") and input_module.tokenizer is not None: return input_module.tokenizer return None @property def max_seq_length(self) -> int | None: # Collect all unique max_seq_length values max_seq_lengths = set() for modules in self.sub_modules.values(): input_module: InputModule = modules[0] if modules and hasattr(input_module, "max_seq_length"): if max_seq_length := input_module.max_seq_length: max_seq_lengths.add(max_seq_length) if not max_seq_lengths: return None elif len(max_seq_lengths) == 1: # Only one unique max_seq_length return max_seq_lengths.pop() else: logger.warning_once(f"Different max_seq_lengths detected: {max_seq_lengths}. Using the maximum value.") return max(max_seq_lengths) @max_seq_length.setter def max_seq_length(self, value) -> None: # Check which modules have max_seq_length has_max_seq_length_keys = [] for key, models in self.sub_modules.items(): if models and hasattr(models[0], "max_seq_length"): has_max_seq_length_keys.append(key) if len(has_max_seq_length_keys) == 0: logger.warning("No modules have a max_seq_length attribute to set.") return for key in has_max_seq_length_keys: input_module: InputModule = self.sub_modules[key][0] input_module.max_seq_length = value # For backwards compatibility, we ensure that the legacy `Asym` alias points to the new `Router` class. Asym = Router