Skip to content

Backbone API

Interfaces and contracts

interfaces

Capability-based interfaces for vision-language backbones.

The papers in :mod:methods do not all consume the same kind of encoder. Some need a black-box text encoder, some inject learnable tokens into a text transformer, and SLDPC needs a paired slide projector and promptable text tower. A single (model, tokenizer) tuple cannot express those differences safely.

This module deliberately keeps the interfaces narrow. Vendor-specific wrappers implement them while method code keeps ownership of the published architecture. In particular, compatibility never implies that an unpublished alignment projection may be inserted between unrelated models.

BackboneCompatibilityError

Bases: ValueError

Raised before model construction when an encoder contract is unmet.

BackboneCapability

Bases: str, Enum

Small, independently testable encoder capabilities.

FeatureLevel

Bases: str, Enum

The tensor level at which a method meets its encoder.

SwapPolicy

Bases: str, Enum

How broadly a method can accept another encoder implementation.

TokenBatch dataclass

TokenBatch(
    input_ids: Tensor,
    attention_mask: Optional[Tensor] = None,
    eot_indices: Optional[Tensor] = None,
    extras: Mapping[str, Tensor] = dict(),
)

Tokenizer output with explicit masks and pooling positions.

extras preserves vendor fields such as token_type_ids without making every consumer depend on a Hugging Face BatchEncoding.

to

to(device: Any) -> 'TokenBatch'

Return a copy with every tensor moved to device.

Source code in common/backbones/interfaces.py
def to(self, device: Any) -> "TokenBatch":
    """Return a copy with every tensor moved to ``device``."""
    return TokenBatch(
        input_ids=self.input_ids.to(device),
        attention_mask=(self.attention_mask.to(device)
                        if self.attention_mask is not None else None),
        eot_indices=(self.eot_indices.to(device)
                     if self.eot_indices is not None else None),
        extras={key: value.to(device) for key, value in self.extras.items()},
    )

as_kwargs

as_kwargs() -> Dict[str, Tensor]

Convert the batch to keyword arguments accepted by text models.

Source code in common/backbones/interfaces.py
def as_kwargs(self) -> Dict[str, torch.Tensor]:
    """Convert the batch to keyword arguments accepted by text models."""
    output = {"input_ids": self.input_ids, **dict(self.extras)}
    if self.attention_mask is not None:
        output["attention_mask"] = self.attention_mask
    return output

BackboneSpec dataclass

BackboneSpec(
    name: str,
    family: str,
    feature_space_id: str,
    capabilities: FrozenSet[BackboneCapability],
    revision: Optional[str] = None,
    tile_dim: Optional[int] = None,
    slide_input_dim: Optional[int] = None,
    vision_token_dim: Optional[int] = None,
    text_token_dim: Optional[int] = None,
    shared_dim: Optional[int] = None,
    context_length: Optional[int] = None,
    image_size: Optional[int] = None,
    aliases: Tuple[str, ...] = (),
)

Record dimensions and provenance needed to judge a backbone swap.

Capabilities describe operations exposed by the runtime bundle. Dimension fields refer to native encoder boundaries; shared_dim is the paired comparison space when one exists.

has

has(*capabilities: BackboneCapability) -> bool

Return whether the spec declares every requested capability.

Source code in common/backbones/interfaces.py
def has(self, *capabilities: BackboneCapability) -> bool:
    """Return whether the spec declares every requested capability."""
    return all(capability in self.capabilities for capability in capabilities)

TextEncoder

Bases: Protocol

Black-box text encoding in the model's native shared space.

tokenize

tokenize(texts: Sequence[str]) -> TokenBatch

Tokenize ordered strings without changing their order.

Source code in common/backbones/interfaces.py
def tokenize(self, texts: Sequence[str]) -> TokenBatch:
    """Tokenize ordered strings without changing their order."""
    ...

encode_text

encode_text(
    texts_or_tokens: Any, normalize: bool = True
) -> Tensor

Encode strings or tokens into the native shared feature space.

Source code in common/backbones/interfaces.py
def encode_text(self, texts_or_tokens: Any,
                normalize: bool = True) -> torch.Tensor:
    """Encode strings or tokens into the native shared feature space."""
    ...

PromptableTextEncoder

Bases: TextEncoder, Protocol

Text tower that supports differentiable embedded soft prompts.

token_width property

token_width: int

Return the width of native token embeddings.

dtype property

dtype: dtype

Return the computation dtype of the promptable text tower.

device property

device: device

Return the device hosting the promptable text tower.

tokenizer property

tokenizer: Any

Return the native tokenizer used by this text tower.

embed_tokens

embed_tokens(tokens: TokenBatch) -> Tensor

Map token IDs to differentiable native token embeddings.

Source code in common/backbones/interfaces.py
def embed_tokens(self, tokens: TokenBatch) -> torch.Tensor:
    """Map token IDs to differentiable native token embeddings."""
    ...

encode_embedded

encode_embedded(
    embeddings: Tensor,
    tokens: TokenBatch,
    normalize: bool = True,
) -> Tensor

Encode pre-embedded prompt tokens into the paired shared space.

Source code in common/backbones/interfaces.py
def encode_embedded(self, embeddings: torch.Tensor, tokens: TokenBatch,
                    normalize: bool = True) -> torch.Tensor:
    """Encode pre-embedded prompt tokens into the paired shared space."""
    ...

tokenize

tokenize(texts: Sequence[str]) -> TokenBatch

Tokenize ordered strings without changing their order.

Source code in common/backbones/interfaces.py
def tokenize(self, texts: Sequence[str]) -> TokenBatch:
    """Tokenize ordered strings without changing their order."""
    ...

encode_text

encode_text(
    texts_or_tokens: Any, normalize: bool = True
) -> Tensor

Encode strings or tokens into the native shared feature space.

Source code in common/backbones/interfaces.py
def encode_text(self, texts_or_tokens: Any,
                normalize: bool = True) -> torch.Tensor:
    """Encode strings or tokens into the native shared feature space."""
    ...

TileEncoder

Bases: Protocol

Black-box image-tile encoder in the model's native feature space.

encode_tiles

encode_tiles(
    images: Tensor, normalize: bool = True
) -> Tensor

Encode a batch of image tiles, optionally normalizing each row.

Source code in common/backbones/interfaces.py
def encode_tiles(self, images: torch.Tensor,
                 normalize: bool = True) -> torch.Tensor:
    """Encode a batch of image tiles, optionally normalizing each row."""
    ...

SlideProjector

Bases: Protocol

Native projection from raw slide vectors to a paired text space.

project_slide

project_slide(
    raw_embeddings: Tensor, normalize: bool = True
) -> Tensor

Project raw slide embeddings through the paired model head.

Source code in common/backbones/interfaces.py
def project_slide(self, raw_embeddings: torch.Tensor,
                  normalize: bool = True) -> torch.Tensor:
    """Project raw slide embeddings through the paired model head."""
    ...

EncoderBundle dataclass

EncoderBundle(
    raw_model: Module,
    spec: BackboneSpec,
    raw_tokenizer: Any = None,
    text: Optional[TextEncoder] = None,
    tile: Optional[TileEncoder] = None,
    slide: Optional[SlideProjector] = None,
    preprocess: Optional[Callable[..., Any]] = None,
    metadata: Mapping[str, Any] = dict(),
)

Uniform runtime handle while retaining access to the native objects.

raw_model and raw_tokenizer are intentional escape hatches for paper implementations that traverse a specific transformer's layers. Capability validation must happen before those objects are passed on.

model property

model: Module

Return the native model (migration alias for raw_model).

tokenizer property

tokenizer: Any

Return the native tokenizer (migration alias for raw_tokenizer).

require

require(
    *capabilities: BackboneCapability,
    consumer: Optional[str] = None,
) -> "EncoderBundle"

Assert that this bundle advertises the requested operations.

Parameters:

Name Type Description Default
*capabilities BackboneCapability

Operations required by the caller.

()
consumer Optional[str]

Optional name included in validation errors.

None

Returns:

Type Description
'EncoderBundle'

This bundle, enabling fluent validation.

Raises:

Type Description
BackboneCompatibilityError

If any capability is absent.

Source code in common/backbones/interfaces.py
def require(self, *capabilities: BackboneCapability,
            consumer: Optional[str] = None) -> "EncoderBundle":
    """Assert that this bundle advertises the requested operations.

    Args:
        *capabilities: Operations required by the caller.
        consumer: Optional name included in validation errors.

    Returns:
        This bundle, enabling fluent validation.

    Raises:
        BackboneCompatibilityError: If any capability is absent.
    """
    missing = [item.value for item in capabilities
               if item not in self.spec.capabilities]
    if missing:
        prefix = f"{consumer} requires" if consumer else "Encoder requires"
        raise BackboneCompatibilityError(
            f"{prefix} {missing}, but backbone '{self.spec.name}' provides "
            f"{sorted(item.value for item in self.spec.capabilities)}.")
    return self

freeze

freeze() -> 'EncoderBundle'

Put the native model in evaluation mode and disable gradients.

Source code in common/backbones/interfaces.py
def freeze(self) -> "EncoderBundle":
    """Put the native model in evaluation mode and disable gradients."""
    self.raw_model.eval()
    self.raw_model.requires_grad_(False)
    return self

encode_text

encode_text(
    texts_or_tokens: Any, normalize: bool = True
) -> Tensor

Encode text using the bundle's validated text wrapper.

Raises:

Type Description
BackboneCompatibilityError

If text encoding is unavailable or the spec declares it without providing a wrapper.

Source code in common/backbones/interfaces.py
def encode_text(self, texts_or_tokens: Any,
                normalize: bool = True) -> torch.Tensor:
    """Encode text using the bundle's validated text wrapper.

    Raises:
        BackboneCompatibilityError: If text encoding is unavailable or
            the spec declares it without providing a wrapper.
    """
    self.require(BackboneCapability.TEXT_ENCODE)
    if self.text is None:
        raise BackboneCompatibilityError(
            f"Backbone '{self.spec.name}' declares text encoding but has no text wrapper.")
    return self.text.encode_text(texts_or_tokens, normalize=normalize)

project_slide

project_slide(
    raw_embeddings: Tensor, normalize: bool = True
) -> Tensor

Apply the bundle's native slide projector.

Raises:

Type Description
BackboneCompatibilityError

If native slide projection is absent.

Source code in common/backbones/interfaces.py
def project_slide(self, raw_embeddings: torch.Tensor,
                  normalize: bool = True) -> torch.Tensor:
    """Apply the bundle's native slide projector.

    Raises:
        BackboneCompatibilityError: If native slide projection is absent.
    """
    self.require(BackboneCapability.SLIDE_PROJECT)
    if self.slide is None:
        raise BackboneCompatibilityError(
            f"Backbone '{self.spec.name}' declares slide projection but has no projector.")
    return self.slide.project_slide(raw_embeddings, normalize=normalize)

assert_feature_space

assert_feature_space(
    *,
    feature_space_id: Optional[str] = None,
    dimension: Optional[int] = None,
    level: FeatureLevel = PATCH_BAG,
) -> None

Validate cached feature provenance and width against the spec.

Parameters:

Name Type Description Default
feature_space_id Optional[str]

Exact producer/checkpoint identity, when known.

None
dimension Optional[int]

Last dimension of the cached feature tensor.

None
level FeatureLevel

Tensor level used to select tile versus slide dimensions.

PATCH_BAG

Raises:

Type Description
BackboneCompatibilityError

If provenance or width differs.

Source code in common/backbones/interfaces.py
def assert_feature_space(self, *, feature_space_id: Optional[str] = None,
                         dimension: Optional[int] = None,
                         level: FeatureLevel = FeatureLevel.PATCH_BAG) -> None:
    """Validate cached feature provenance and width against the spec.

    Args:
        feature_space_id: Exact producer/checkpoint identity, when known.
        dimension: Last dimension of the cached feature tensor.
        level: Tensor level used to select tile versus slide dimensions.

    Raises:
        BackboneCompatibilityError: If provenance or width differs.
    """
    if feature_space_id and feature_space_id != self.spec.feature_space_id:
        raise BackboneCompatibilityError(
            f"Feature space '{feature_space_id}' does not match backbone "
            f"'{self.spec.name}' ({self.spec.feature_space_id}).")
    native_dim = (self.spec.slide_input_dim
                  if level == FeatureLevel.SLIDE_EMBEDDING
                  else self.spec.tile_dim)
    if dimension is not None and native_dim is not None and dimension != native_dim:
        raise BackboneCompatibilityError(
            f"{level.value} width {dimension} does not match backbone "
            f"'{self.spec.name}' width {native_dim}.")

MethodBackboneContract dataclass

MethodBackboneContract(
    method: str,
    feature_level: FeatureLevel,
    swap_policy: SwapPolicy,
    required_capabilities: FrozenSet[
        BackboneCapability
    ] = frozenset(),
    config_key: Optional[str] = "backbone",
    default_backbone: Optional[str] = None,
    supported_backbones: Tuple[str, ...] = (),
    name_aliases: Mapping[str, str] = dict(),
    feature_dims: Mapping[str, Tuple[int, ...]] = dict(),
    feature_boundaries: Mapping[
        str, Mapping[str, Tuple[int, ...]]
    ] = dict(),
    feature_dim_key: Optional[str] = "feature_dim",
    feature_space_key: Optional[str] = "feature_space_id",
    bundle_feature_space_key: Optional[str] = None,
    enforce_native_dimension: bool = False,
    require_feature_space: bool = False,
    rationale: str = "",
)

Declare the encoder and feature boundary accepted by a method.

swap_policy distinguishes architectural compatibility from mere tensor shape compatibility. Allowlisted, fixed, and precomputed methods must name a supported native family; capability-based methods validate operations on the returned bundle.

resolve_name

resolve_name(cfg: Mapping[str, Any]) -> Optional[str]

Resolve and canonicalize the backbone selected by a run config.

Source code in common/backbones/interfaces.py
def resolve_name(self, cfg: Mapping[str, Any]) -> Optional[str]:
    """Resolve and canonicalize the backbone selected by a run config."""
    if self.config_key is None:
        return self.default_backbone
    value = cfg.get(self.config_key, self.default_backbone)
    if value is None:
        return None
    raw = str(value).strip().lower().replace("_", "-")
    return canonical_backbone_name(self.name_aliases.get(raw, raw))

validate_config

validate_config(cfg: Mapping[str, Any]) -> Optional[str]

Validate config-only constraints before allocating model weights.

Returns:

Type Description
Optional[str]

The canonical selected backbone, or None for a contract with

Optional[str]

no runtime encoder name.

Raises:

Type Description
BackboneCompatibilityError

If the name, width, provenance, or required config fields violate this contract.

Source code in common/backbones/interfaces.py
def validate_config(self, cfg: Mapping[str, Any]) -> Optional[str]:
    """Validate config-only constraints before allocating model weights.

    Returns:
        The canonical selected backbone, or ``None`` for a contract with
        no runtime encoder name.

    Raises:
        BackboneCompatibilityError: If the name, width, provenance, or
            required config fields violate this contract.
    """
    selected = self.resolve_name(cfg)
    allowed = tuple(canonical_backbone_name(item)
                    for item in self.supported_backbones)
    if self.swap_policy in {SwapPolicy.ALLOWLIST, SwapPolicy.FIXED,
                            SwapPolicy.PRECOMPUTED}:
        if selected is None or selected not in allowed:
            raise BackboneCompatibilityError(
                f"Method '{self.method}' cannot use backbone {selected!r}. "
                f"Supported: {list(allowed)}. {self.rationale}".strip())
    if (self.require_feature_space and self.feature_space_key and
            not cfg.get(self.feature_space_key)):
        raise BackboneCompatibilityError(
            f"Method '{self.method}' requires '{self.feature_space_key}' "
            "to identify the offline encoder checkpoint exactly.")
    if (self.enforce_native_dimension and self.feature_dim_key and
            self.feature_dim_key not in cfg):
        raise BackboneCompatibilityError(
            f"Method '{self.method}' requires '{self.feature_dim_key}' "
            f"for its {self.feature_level.value} input.")
    if self.feature_dim_key and self.feature_dim_key in cfg and selected is not None:
        expected = self.feature_dims.get(selected)
        actual = int(cfg[self.feature_dim_key])
        if expected and actual not in expected:
            raise BackboneCompatibilityError(
                f"Method '{self.method}' expects {selected} "
                f"{self.feature_level.value} width in {list(expected)}, got {actual}.")
    boundaries = self.feature_boundaries.get(selected or "", {})
    if boundaries and self.feature_space_key and cfg.get(self.feature_space_key):
        feature_space = str(cfg[self.feature_space_key])
        expected_dims = boundaries.get(feature_space)
        if expected_dims is None:
            raise BackboneCompatibilityError(
                f"Feature space mismatch: method '{self.method}' cannot "
                "consume feature boundary "
                f"{feature_space!r} for backbone {selected!r}. Supported: "
                f"{sorted(boundaries)}.")
        if self.feature_dim_key and self.feature_dim_key in cfg:
            actual = int(cfg[self.feature_dim_key])
            if actual not in expected_dims:
                raise BackboneCompatibilityError(
                    f"Feature boundary {feature_space!r} requires width in "
                    f"{list(expected_dims)}, got {actual}.")
    return selected

validate_bundle

validate_bundle(
    cfg: Mapping[str, Any], bundle: EncoderBundle
) -> EncoderBundle

Validate a loaded bundle against this contract and run config.

Returns:

Type Description
EncoderBundle

The validated bundle.

Raises:

Type Description
BackboneCompatibilityError

If capabilities, identity, dimensions, or feature provenance are incompatible.

Source code in common/backbones/interfaces.py
def validate_bundle(self, cfg: Mapping[str, Any],
                    bundle: EncoderBundle) -> EncoderBundle:
    """Validate a loaded bundle against this contract and run config.

    Returns:
        The validated bundle.

    Raises:
        BackboneCompatibilityError: If capabilities, identity, dimensions,
            or feature provenance are incompatible.
    """
    selected = self.validate_config(cfg)
    if self.swap_policy != SwapPolicy.CAPABILITY and selected != bundle.spec.name:
        raise BackboneCompatibilityError(
            f"Method '{self.method}' selected '{selected}', but loader returned "
            f"'{bundle.spec.name}'.")
    bundle.require(*self.required_capabilities, consumer=f"Method '{self.method}'")
    bundle_feature_space = (
        cfg.get(self.bundle_feature_space_key)
        if self.bundle_feature_space_key else None)
    if self.bundle_feature_space_key and not bundle_feature_space:
        raise BackboneCompatibilityError(
            f"Method '{self.method}' requires "
            f"'{self.bundle_feature_space_key}' to identify its runtime "
            "encoder checkpoint exactly.")
    if bundle_feature_space:
        bundle.assert_feature_space(
            feature_space_id=str(bundle_feature_space))
    feature_space = (cfg.get(self.feature_space_key)
                     if self.feature_space_key else None)
    if self.require_feature_space and not feature_space:
        raise BackboneCompatibilityError(
            f"Method '{self.method}' requires '{self.feature_space_key}' "
            "to identify the offline encoder checkpoint exactly.")
    dimension = (int(cfg[self.feature_dim_key])
                 if self.feature_dim_key and self.feature_dim_key in cfg else None)
    # A learned input adapter is part of some native methods (e.g. MUSE),
    # so only enforce native width when this contract declares exact dims.
    boundaries = self.feature_boundaries.get(selected or "", {})
    if boundaries:
        # Some methods explicitly restore the encoder's shared space from
        # a named internal boundary (for example a frozen CLIP visual
        # projection applied to cached vision-preprojection tensors).
        # validate_config() has already checked the exact source identity
        # and width; the runtime bundle check here still proves that the
        # matching encoder family and capabilities were loaded.
        # Most boundary-producing vision towers belong to the runtime
        # encoder itself. A paired method can instead identify and check
        # a distinct runtime text/slide encoder explicitly.
        if not self.bundle_feature_space_key:
            accepted = tuple(boundaries)
            if not all(
                    value == bundle.spec.feature_space_id
                    or value.startswith(bundle.spec.feature_space_id + "#")
                    for value in accepted):
                raise BackboneCompatibilityError(
                    f"Method '{self.method}' declares feature boundaries "
                    f"that do not belong to backbone "
                    f"'{bundle.spec.name}'.")
    elif selected in self.feature_dims or self.enforce_native_dimension:
        bundle.assert_feature_space(feature_space_id=feature_space,
                                    dimension=dimension,
                                    level=self.feature_level)
    elif feature_space:
        bundle.assert_feature_space(feature_space_id=feature_space,
                                    level=self.feature_level)
    return bundle

as_dict

as_dict() -> Dict[str, Any]

Serialize the contract to JSON-compatible registry metadata.

Source code in common/backbones/interfaces.py
def as_dict(self) -> Dict[str, Any]:
    """Serialize the contract to JSON-compatible registry metadata."""
    return {
        "method": self.method,
        "feature_level": self.feature_level.value,
        "swap_policy": self.swap_policy.value,
        "config_key": self.config_key,
        "default_backbone": self.default_backbone,
        "supported_backbones": list(self.supported_backbones),
        "name_aliases": dict(self.name_aliases),
        "required_capabilities": sorted(item.value for item in self.required_capabilities),
        "feature_dims": {key: list(value) for key, value in self.feature_dims.items()},
        "feature_boundaries": {
            backbone: {
                feature_space: list(dimensions)
                for feature_space, dimensions in boundaries.items()
            }
            for backbone, boundaries in self.feature_boundaries.items()
        },
        "bundle_feature_space_key": self.bundle_feature_space_key,
        "enforce_native_dimension": self.enforce_native_dimension,
        "require_feature_space": self.require_feature_space,
        "rationale": self.rationale,
    }

canonical_backbone_name

canonical_backbone_name(name: Any) -> str

Normalize config spellings without erasing meaningful model families.

Source code in common/backbone_names.py
def canonical_backbone_name(name: Any) -> str:
    """Normalize config spellings without erasing meaningful model families."""
    value = str(name).strip().lower().replace("_", "-")
    return _ALIASES.get(value, value)

normalize_features

normalize_features(
    features: Tensor, normalize: bool
) -> Tensor

Convert features to float and optionally L2-normalize the last axis.

Source code in common/backbones/interfaces.py
def normalize_features(features: torch.Tensor, normalize: bool) -> torch.Tensor:
    """Convert features to float and optionally L2-normalize the last axis."""
    return F.normalize(features.float(), dim=-1) if normalize else features

Registry and builders

factory

Backbone registry and vendor adapters.

build_backbone retains the historical (model, tokenizer, info) return value. New code should use :func:build_encoder, whose :class:~common.backbones.interfaces.EncoderBundle exposes explicit capabilities and feature-space provenance.

BackboneInfo dataclass

BackboneInfo(
    name: str,
    patch_dim: int,
    text_dim: int,
    image_size: int = 224,
    is_clip_compatible: bool = True,
)

Backward-compatible metadata returned by build_backbone.

patch_dim and text_dim are the projected comparison-space dimensions. Use :class:BackboneSpec when token widths or raw slide dimensions matter.

TitanPromptableText

TitanPromptableText(titan: Module)

Expose TITAN's text tower through the promptable-text protocol.

The wrapper preserves TITAN's native tokenizer, end-of-text pooling, token width, dtype, and paired projection. SLDPC uses it to optimize context embeddings without replacing the native text tower.

Source code in common/backbones/factory.py
def __init__(self, titan: nn.Module):
    self.model = titan
    self._text = titan.text_encoder

list_backbones

list_backbones() -> list[str]

Return canonical encoder names in registry order.

Source code in common/backbones/factory.py
def list_backbones() -> list[str]:
    """Return canonical encoder names in registry order."""
    return list(_SPECS)

get_spec

get_spec(name: str) -> BackboneSpec

Return immutable capability and provenance metadata for an encoder.

Parameters:

Name Type Description Default
name str

Canonical encoder name or registered alias.

required

Raises:

Type Description
KeyError

If no matching encoder is registered.

Source code in common/backbones/factory.py
def get_spec(name: str) -> BackboneSpec:
    """Return immutable capability and provenance metadata for an encoder.

    Args:
        name: Canonical encoder name or registered alias.

    Raises:
        KeyError: If no matching encoder is registered.
    """
    canonical = canonical_backbone_name(name)
    if canonical not in _SPECS:
        raise KeyError(f"Unknown backbone '{name}'. Available: {list_backbones()}")
    return _SPECS[canonical]

get_info

get_info(name: str) -> BackboneInfo

Return legacy projected dimensions for a registered encoder.

Source code in common/backbones/factory.py
def get_info(name: str) -> BackboneInfo:
    """Return legacy projected dimensions for a registered encoder."""
    return _legacy_info(get_spec(name))

register_backbone

register_backbone(
    spec: BackboneSpec,
    builder: EncoderBuilder,
    *,
    overwrite: bool = False,
) -> None

Register an encoder bundle without modifying any method architecture.

A builder receives the same keyword arguments as :func:build_encoder and must return a bundle whose spec.name matches spec.name.

Parameters:

Name Type Description Default
spec BackboneSpec

Immutable metadata for the new encoder.

required
builder EncoderBuilder

Callable returning a compatible :class:EncoderBundle.

required
overwrite bool

Whether to replace an existing registration deliberately.

False

Raises:

Type Description
KeyError

If the name already exists and overwrite is false.

ValueError

If spec.name is not already canonical.

Source code in common/backbones/factory.py
def register_backbone(spec: BackboneSpec, builder: EncoderBuilder,
                      *, overwrite: bool = False) -> None:
    """Register an encoder bundle without modifying any method architecture.

    A builder receives the same keyword arguments as :func:`build_encoder`
    and must return a bundle whose ``spec.name`` matches ``spec.name``.

    Args:
        spec: Immutable metadata for the new encoder.
        builder: Callable returning a compatible :class:`EncoderBundle`.
        overwrite: Whether to replace an existing registration deliberately.

    Raises:
        KeyError: If the name already exists and ``overwrite`` is false.
        ValueError: If ``spec.name`` is not already canonical.
    """
    name = canonical_backbone_name(spec.name)
    if name in _SPECS and not overwrite:
        raise KeyError(f"Backbone '{name}' is already registered")
    if name != spec.name:
        raise ValueError("BackboneSpec.name must already be canonical")
    _SPECS[name] = spec
    _CUSTOM_BUILDERS[name] = builder

unregister_backbone

unregister_backbone(name: str) -> None

Remove a process-local custom registration.

Raises:

Type Description
KeyError

If name is built in or is not currently registered as a custom encoder.

Source code in common/backbones/factory.py
def unregister_backbone(name: str) -> None:
    """Remove a process-local custom registration.

    Raises:
        KeyError: If ``name`` is built in or is not currently registered as a
            custom encoder.
    """
    canonical = canonical_backbone_name(name)
    if canonical not in _CUSTOM_BUILDERS:
        raise KeyError(f"'{name}' is not a custom backbone")
    _CUSTOM_BUILDERS.pop(canonical)
    _SPECS.pop(canonical)

build_encoder

build_encoder(
    name: str,
    weights_path: Optional[str] = None,
    device: str = "cuda",
    **kwargs: Any,
) -> EncoderBundle

Load a capability-aware encoder bundle.

Parameters:

Name Type Description Default
name str

Canonical encoder name or alias.

required
weights_path Optional[str]

Optional local checkpoint, snapshot, or model identifier.

None
device str

PyTorch device on which to construct the native model.

'cuda'
**kwargs Any

Family-specific loader options. TITAN accepts model_id, revision, and local_files_only.

{}

Returns:

Type Description
EncoderBundle

The native model, tokenizer, provenance spec, and supported wrappers.

Raises:

Type Description
KeyError

If the encoder is unknown.

TypeError

If loader options are unsupported for the selected family.

BackboneCompatibilityError

If a custom builder returns the wrong spec.

Source code in common/backbones/factory.py
def build_encoder(name: str, weights_path: Optional[str] = None,
                  device: str = "cuda", **kwargs: Any) -> EncoderBundle:
    """Load a capability-aware encoder bundle.

    Args:
        name: Canonical encoder name or alias.
        weights_path: Optional local checkpoint, snapshot, or model identifier.
        device: PyTorch device on which to construct the native model.
        **kwargs: Family-specific loader options. TITAN accepts ``model_id``,
            ``revision``, and ``local_files_only``.

    Returns:
        The native model, tokenizer, provenance spec, and supported wrappers.

    Raises:
        KeyError: If the encoder is unknown.
        TypeError: If loader options are unsupported for the selected family.
        BackboneCompatibilityError: If a custom builder returns the wrong spec.
    """
    canonical = canonical_backbone_name(name)
    spec = get_spec(canonical)
    if canonical in _CUSTOM_BUILDERS:
        bundle = _CUSTOM_BUILDERS[canonical](
            weights_path=weights_path, device=device, **kwargs)
        if bundle.spec.name != spec.name:
            raise BackboneCompatibilityError(
                f"Builder registered as '{spec.name}' returned '{bundle.spec.name}'.")
        return bundle
    if canonical == "titan":
        model_id = kwargs.pop(
            "model_id", weights_path or "MahmoodLab/TITAN")
        revision = kwargs.pop("revision", None)
        local_files_only = kwargs.pop("local_files_only", False)
        if kwargs:
            unknown = ", ".join(sorted(kwargs))
            raise TypeError(
                f"Unexpected loader options for '{canonical}': {unknown}")
        return _load_titan(
            model_id, device, revision=revision,
            local_files_only=local_files_only)
    if kwargs:
        unknown = ", ".join(sorted(kwargs))
        raise TypeError(f"Unexpected loader options for '{canonical}': {unknown}")
    model, tokenizer, preprocess = _load_builtin(canonical, weights_path, device)
    if spec.family == "openai_clip":
        text = _OpenAIClipText(model, tokenizer)
    elif spec.family == "hf_clip":
        text = _HFClipText(model, tokenizer)
    elif spec.family == "conch":
        text = _ConchText(model, tokenizer)
    elif spec.family == "musk":
        text = _MuskText(model, tokenizer, max_len=spec.context_length or 100)
    else:
        text = _NativeText(model, tokenizer)
    return EncoderBundle(model, spec, raw_tokenizer=tokenizer,
                         text=text, tile=_NativeTile(model), preprocess=preprocess)

build_backbone

build_backbone(
    name: str,
    weights_path: Optional[str] = None,
    device: str = "cuda",
) -> Tuple[Module, Callable, BackboneInfo]

Load the legacy (model, tokenizer, info) tuple.

New framework integrations should call :func:build_encoder; this wrapper exists for vendored implementations that still consume native objects.

Source code in common/backbones/factory.py
def build_backbone(name: str, weights_path: Optional[str] = None,
                   device: str = "cuda") -> Tuple[nn.Module, Callable, BackboneInfo]:
    """Load the legacy ``(model, tokenizer, info)`` tuple.

    New framework integrations should call :func:`build_encoder`; this wrapper
    exists for vendored implementations that still consume native objects.
    """
    bundle = build_encoder(name, weights_path=weights_path, device=device)
    return bundle.raw_model, bundle.raw_tokenizer, _legacy_info(bundle.spec)