spatial-building-embeddings / modeling_spatial_embeddings.py

fix: set all_tied_weights_keys for transformers 5.6 compatibility

7ebdef3 verified 6 months ago

7.97 kB

from typing import Literal, Optional, Tuple, Union

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import AutoConfig, AutoModel, PreTrainedModel

# Handle import for both local development and HuggingFace Hub loading
try:
    from .configuration_spatial_embeddings import SpatialEmbeddingsConfig
except ImportError:
    # When loaded from HuggingFace Hub, relative imports may not work
    # Try absolute import instead
    try:
        from configuration_spatial_embeddings import SpatialEmbeddingsConfig
    except ImportError:
        # Last resort: import from the module directly (Hugging Face dynamic module loading)
        from configuration_spatial_embeddings import SpatialEmbeddingsConfig


class EmbeddingProjector(nn.Module):
    """
    Configurable MLP projection head for embedding transformation.
    (Copied from embedding_pipeline/train/model.py for self-contained publishing)
    """

    def __init__(
        self,
        input_dim: int = 768,
        hidden_dim: int = 512,
        output_dim: int = 256,
        dropout: float = 0.1,
        num_hidden_layers: int = 1,
        hidden_dim_multiplier: float = 1.0,
        activation: Literal["gelu", "relu", "silu"] = "gelu",
        use_residual: bool = True,
        use_layer_norm: bool = True,
    ):
        super().__init__()

        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        self.output_dim = output_dim
        self.use_residual = use_residual
        self.use_layer_norm = use_layer_norm
        self.num_hidden_layers = num_hidden_layers
        self.hidden_dim_multiplier = hidden_dim_multiplier
        self.activation_name = activation

        self.hidden_dims = self._compute_hidden_dims(
            hidden_dim, num_hidden_layers, hidden_dim_multiplier
        )

        self.activation = self._resolve_activation(activation)

        # First hidden block
        self.input_layer = nn.Linear(input_dim, self.hidden_dims[0])
        if use_layer_norm:
            self.input_norm = nn.LayerNorm(self.hidden_dims[0])
        self.input_dropout = nn.Dropout(dropout)

        # Additional hidden blocks (if any)
        self.hidden_layers = nn.ModuleList()
        if use_layer_norm:
            self.hidden_norms = nn.ModuleList()
        else:
            self.hidden_norms = None
        self.hidden_dropouts = nn.ModuleList()

        for idx in range(1, len(self.hidden_dims)):
            layer = nn.Linear(self.hidden_dims[idx - 1], self.hidden_dims[idx])
            self.hidden_layers.append(layer)
            if use_layer_norm:
                self.hidden_norms.append(nn.LayerNorm(self.hidden_dims[idx]))
            self.hidden_dropouts.append(nn.Dropout(dropout))

        # Output block
        self.output_layer = nn.Linear(self.hidden_dims[-1], output_dim)
        if use_layer_norm:
            self.output_norm = nn.LayerNorm(output_dim)
        self.output_dropout = nn.Dropout(dropout)

        # Residual shortcut (projects input directly to output)
        if use_residual:
            self.residual_proj = nn.Linear(input_dim, output_dim)

    @staticmethod
    def _compute_hidden_dims(
        base_hidden_dim: int, num_layers: int, multiplier: float
    ) -> list[int]:
        dims: list[int] = []
        current_dim = base_hidden_dim
        for layer_idx in range(num_layers):
            if layer_idx == 0:
                dims.append(base_hidden_dim)
            else:
                current_dim = max(16, int(round(current_dim * multiplier)))
                dims.append(current_dim)
        return dims

    @staticmethod
    def _resolve_activation(name: str) -> nn.Module:
        if name == "gelu":
            return nn.GELU()
        if name == "relu":
            return nn.ReLU()
        if name == "silu":
            return nn.SiLU()
        raise ValueError(f"Unsupported activation: {name}")

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # First hidden block
        out = self.input_layer(x)
        if self.use_layer_norm:
            out = self.input_norm(out)
        out = self.activation(out)
        out = self.input_dropout(out)

        # Additional hidden blocks
        for idx, layer in enumerate(self.hidden_layers):
            out = layer(out)
            if self.use_layer_norm and self.hidden_norms is not None:
                out = self.hidden_norms[idx](out)
            out = self.activation(out)
            out = self.hidden_dropouts[idx](out)

        # Output block
        out = self.output_layer(out)
        if self.use_layer_norm:
            out = self.output_norm(out)
        out = self.output_dropout(out)

        # Residual connection
        if self.use_residual:
            residual = self.residual_proj(x)
            out = out + residual

        # L2 normalization
        out = F.normalize(out, p=2, dim=1)

        return out


class SpatialEmbeddingsModel(PreTrainedModel):
    config_class = SpatialEmbeddingsConfig
    _tied_weights_keys = []

    def __init__(self, config: SpatialEmbeddingsConfig):
        super().__init__(config)
        self.config = config
        # transformers ≥5.6 accesses all_tied_weights_keys before tie_weights() is called
        # for custom trust_remote_code models; set it here since we have no tied weights.
        if not hasattr(self, "all_tied_weights_keys"):
            self.all_tied_weights_keys = {}

        # Initialize backbone from config rather than from_pretrained.
        # Calling from_pretrained inside __init__ is an anti-pattern in transformers ≥5:
        # the outer from_pretrained wraps __init__ in a meta-device context
        # (low_cpu_mem_usage=True), and a nested from_pretrained inside that context
        # raises an error. Using from_config is safe in any device context.
        # Backbone weights are populated by the outer from_pretrained weight-loading
        # step (inference) or explicitly by convert_to_hf.py (publishing).
        backbone_config = AutoConfig.from_pretrained(
            config.backbone_model_name, trust_remote_code=True
        )
        self.backbone = AutoModel.from_config(backbone_config, trust_remote_code=True)

        # Initialize projector
        self.projector = EmbeddingProjector(
            input_dim=config.input_dim,
            hidden_dim=config.hidden_dim,
            output_dim=config.output_dim,
            dropout=config.dropout,
            num_hidden_layers=config.num_hidden_layers,
            hidden_dim_multiplier=config.hidden_dim_multiplier,
            activation=config.activation,
            use_residual=config.use_residual,
            use_layer_norm=config.use_layer_norm,
        )

    def forward(
        self,
        pixel_values: Optional[torch.Tensor] = None,
        return_dict: Optional[bool] = None,
        **kwargs,
    ) -> Union[Tuple, torch.Tensor]:
        """
        Args:
            pixel_values: Tensor of shape (batch_size, channels, height, width)
            return_dict: Whether to return a dictionary or tuple
        Returns:
            If return_dict is True (default for HF), returns object with 'embeddings'.
            Otherwise returns (embeddings,).
        """
        # Pass through backbone
        outputs = self.backbone(pixel_values=pixel_values, return_dict=True, **kwargs)

        # Extract pooled output (CLS token or similar)
        # DINOv2 outputs pooler_output in some versions, or last_hidden_state
        if hasattr(outputs, "pooler_output") and outputs.pooler_output is not None:
            backbone_emb = outputs.pooler_output
        else:
            # Fallback: Use CLS token from last hidden state
            backbone_emb = outputs.last_hidden_state[:, 0]

        # Project to specialized embedding
        specialized_emb = self.projector(backbone_emb)

        if return_dict:
            return {"embeddings": specialized_emb, "backbone_outputs": outputs}
        return (specialized_emb,)

	from typing import Literal, Optional, Tuple, Union

	import torch
	import torch.nn as nn
	import torch.nn.functional as F
	from transformers import AutoConfig, AutoModel, PreTrainedModel

	# Handle import for both local development and HuggingFace Hub loading
	try:
	from .configuration_spatial_embeddings import SpatialEmbeddingsConfig
	except ImportError:
	# When loaded from HuggingFace Hub, relative imports may not work
	# Try absolute import instead
	try:
	from configuration_spatial_embeddings import SpatialEmbeddingsConfig
	except ImportError:
	# Last resort: import from the module directly (Hugging Face dynamic module loading)
	from configuration_spatial_embeddings import SpatialEmbeddingsConfig


	class EmbeddingProjector(nn.Module):
	"""
	Configurable MLP projection head for embedding transformation.
	(Copied from embedding_pipeline/train/model.py for self-contained publishing)
	"""

	def __init__(
	self,
	input_dim: int = 768,
	hidden_dim: int = 512,
	output_dim: int = 256,
	dropout: float = 0.1,
	num_hidden_layers: int = 1,
	hidden_dim_multiplier: float = 1.0,
	activation: Literal["gelu", "relu", "silu"] = "gelu",
	use_residual: bool = True,
	use_layer_norm: bool = True,
	):
	super().__init__()

	self.input_dim = input_dim
	self.hidden_dim = hidden_dim
	self.output_dim = output_dim
	self.use_residual = use_residual
	self.use_layer_norm = use_layer_norm
	self.num_hidden_layers = num_hidden_layers
	self.hidden_dim_multiplier = hidden_dim_multiplier
	self.activation_name = activation

	self.hidden_dims = self._compute_hidden_dims(
	hidden_dim, num_hidden_layers, hidden_dim_multiplier
	)

	self.activation = self._resolve_activation(activation)

	# First hidden block
	self.input_layer = nn.Linear(input_dim, self.hidden_dims[0])
	if use_layer_norm:
	self.input_norm = nn.LayerNorm(self.hidden_dims[0])
	self.input_dropout = nn.Dropout(dropout)

	# Additional hidden blocks (if any)
	self.hidden_layers = nn.ModuleList()
	if use_layer_norm:
	self.hidden_norms = nn.ModuleList()
	else:
	self.hidden_norms = None
	self.hidden_dropouts = nn.ModuleList()

	for idx in range(1, len(self.hidden_dims)):
	layer = nn.Linear(self.hidden_dims[idx - 1], self.hidden_dims[idx])
	self.hidden_layers.append(layer)
	if use_layer_norm:
	self.hidden_norms.append(nn.LayerNorm(self.hidden_dims[idx]))
	self.hidden_dropouts.append(nn.Dropout(dropout))

	# Output block
	self.output_layer = nn.Linear(self.hidden_dims[-1], output_dim)
	if use_layer_norm:
	self.output_norm = nn.LayerNorm(output_dim)
	self.output_dropout = nn.Dropout(dropout)

	# Residual shortcut (projects input directly to output)
	if use_residual:
	self.residual_proj = nn.Linear(input_dim, output_dim)

	@staticmethod
	def _compute_hidden_dims(
	base_hidden_dim: int, num_layers: int, multiplier: float
	) -> list[int]:
	dims: list[int] = []
	current_dim = base_hidden_dim
	for layer_idx in range(num_layers):
	if layer_idx == 0:
	dims.append(base_hidden_dim)
	else:
	current_dim = max(16, int(round(current_dim * multiplier)))
	dims.append(current_dim)
	return dims

	@staticmethod
	def _resolve_activation(name: str) -> nn.Module:
	if name == "gelu":
	return nn.GELU()
	if name == "relu":
	return nn.ReLU()
	if name == "silu":
	return nn.SiLU()
	raise ValueError(f"Unsupported activation: {name}")

	def forward(self, x: torch.Tensor) -> torch.Tensor:
	# First hidden block
	out = self.input_layer(x)
	if self.use_layer_norm:
	out = self.input_norm(out)
	out = self.activation(out)
	out = self.input_dropout(out)

	# Additional hidden blocks
	for idx, layer in enumerate(self.hidden_layers):
	out = layer(out)
	if self.use_layer_norm and self.hidden_norms is not None:
	out = self.hidden_norms[idx](out)
	out = self.activation(out)
	out = self.hidden_dropouts[idx](out)

	# Output block
	out = self.output_layer(out)
	if self.use_layer_norm:
	out = self.output_norm(out)
	out = self.output_dropout(out)

	# Residual connection
	if self.use_residual:
	residual = self.residual_proj(x)
	out = out + residual

	# L2 normalization
	out = F.normalize(out, p=2, dim=1)

	return out


	class SpatialEmbeddingsModel(PreTrainedModel):
	config_class = SpatialEmbeddingsConfig
	_tied_weights_keys = []

	def __init__(self, config: SpatialEmbeddingsConfig):
	super().__init__(config)
	self.config = config
	# transformers ≥5.6 accesses all_tied_weights_keys before tie_weights() is called
	# for custom trust_remote_code models; set it here since we have no tied weights.
	if not hasattr(self, "all_tied_weights_keys"):
	self.all_tied_weights_keys = {}

	# Initialize backbone from config rather than from_pretrained.
	# Calling from_pretrained inside __init__ is an anti-pattern in transformers ≥5:
	# the outer from_pretrained wraps __init__ in a meta-device context
	# (low_cpu_mem_usage=True), and a nested from_pretrained inside that context
	# raises an error. Using from_config is safe in any device context.
	# Backbone weights are populated by the outer from_pretrained weight-loading
	# step (inference) or explicitly by convert_to_hf.py (publishing).
	backbone_config = AutoConfig.from_pretrained(
	config.backbone_model_name, trust_remote_code=True
	)
	self.backbone = AutoModel.from_config(backbone_config, trust_remote_code=True)

	# Initialize projector
	self.projector = EmbeddingProjector(
	input_dim=config.input_dim,
	hidden_dim=config.hidden_dim,
	output_dim=config.output_dim,
	dropout=config.dropout,
	num_hidden_layers=config.num_hidden_layers,
	hidden_dim_multiplier=config.hidden_dim_multiplier,
	activation=config.activation,
	use_residual=config.use_residual,
	use_layer_norm=config.use_layer_norm,
	)

	def forward(
	self,
	pixel_values: Optional[torch.Tensor] = None,
	return_dict: Optional[bool] = None,
	**kwargs,
	) -> Union[Tuple, torch.Tensor]:
	"""
	Args:
	pixel_values: Tensor of shape (batch_size, channels, height, width)
	return_dict: Whether to return a dictionary or tuple

	Returns:
	If return_dict is True (default for HF), returns object with 'embeddings'.
	Otherwise returns (embeddings,).
	"""
	# Pass through backbone
	outputs = self.backbone(pixel_values=pixel_values, return_dict=True, **kwargs)

	# Extract pooled output (CLS token or similar)
	# DINOv2 outputs pooler_output in some versions, or last_hidden_state
	if hasattr(outputs, "pooler_output") and outputs.pooler_output is not None:
	backbone_emb = outputs.pooler_output
	else:
	# Fallback: Use CLS token from last hidden state
	backbone_emb = outputs.last_hidden_state[:, 0]

	# Project to specialized embedding
	specialized_emb = self.projector(backbone_emb)

	if return_dict:
	return {"embeddings": specialized_emb, "backbone_outputs": outputs}
	return (specialized_emb,)