from dataclasses import dataclass from pathlib import Path import torch from torch import nn import torch.nn.functional as F from huggingface_hub import hf_hub_download from safetensors.torch import load_file from transformers import AutoModel, PreTrainedModel from .configuration_samer import SaMerConfig @dataclass class MergeConfig: num_regions: int = 64 cluster_iters: int = 3 spatial_weight: float = 0.1 assignment_temperature: float = 0.07 image_token_count: int = 1024 class SaMerModel(PreTrainedModel): """Adapter-only SaMer wrapper for ColPali-style multi-vector retrievers.""" config_class = SaMerConfig base_model_prefix = "retriever" _keys_to_ignore_on_load_missing = [r"retriever\..*"] def __init__(self, config: SaMerConfig, **kwargs): super().__init__(config) base_kwargs = {"trust_remote_code": True} local_files_only = kwargs.pop("local_files_only", None) if local_files_only is not None: base_kwargs["local_files_only"] = local_files_only if config.base_model_revision is not None: base_kwargs["revision"] = config.base_model_revision dtype_name = kwargs.pop("dtype", None) or kwargs.pop("torch_dtype", None) or config.base_model_dtype if isinstance(dtype_name, str) and dtype_name: base_kwargs["dtype"] = getattr(torch, dtype_name) elif dtype_name is not None: base_kwargs["dtype"] = dtype_name try: self.retriever = AutoModel.from_pretrained(config.base_model_name_or_path, **base_kwargs) except ValueError as exc: message = str(exc) if "ColPaliConfig" in message: from transformers import ColPaliForRetrieval self.retriever = ColPaliForRetrieval.from_pretrained(config.base_model_name_or_path, **base_kwargs) elif "ColQwen2Config" in message: from transformers import ColQwen2ForRetrieval self.retriever = ColQwen2ForRetrieval.from_pretrained(config.base_model_name_or_path, **base_kwargs) else: raise self.merge_config = MergeConfig( num_regions=config.num_regions, cluster_iters=config.cluster_iters, spatial_weight=config.spatial_weight, assignment_temperature=config.assignment_temperature, image_token_count=config.image_token_count, ) @classmethod def from_pretrained(cls, pretrained_model_name_or_path, *model_args, config=None, **kwargs): """Load the base model plus the SaMer projector without meta-device nesting.""" cache_dir = kwargs.pop("cache_dir", None) force_download = kwargs.pop("force_download", False) local_files_only = kwargs.pop("local_files_only", False) revision = kwargs.pop("revision", None) token = kwargs.pop("token", None) kwargs.pop("trust_remote_code", None) dtype = kwargs.pop("dtype", kwargs.pop("torch_dtype", None)) if config is None: config = SaMerConfig.from_pretrained( pretrained_model_name_or_path, cache_dir=cache_dir, force_download=force_download, local_files_only=local_files_only, revision=revision, token=token, trust_remote_code=True, ) model = cls(config, dtype=dtype, local_files_only=local_files_only) weights_path = _resolve_weights_path( pretrained_model_name_or_path, cache_dir=cache_dir, force_download=force_download, local_files_only=local_files_only, revision=revision, token=token, ) tensors = load_file(weights_path) projector = _find_projector(model.base_model) projector.load_state_dict( { "weight": tensors["base_model.embedding_proj_layer.weight"], "bias": tensors["base_model.embedding_proj_layer.bias"], }, strict=True, ) return model def forward(self, *args, **kwargs): """Run the underlying ColPali/ColQwen2 model without token merging.""" return self.retriever(*args, **kwargs) @torch.inference_mode() def encode_image(self, inputs=None, coords=None, **kwargs): """Encode images and return SaMer-compressed image-side tokens. Args: inputs: Optional processor output dictionary. coords: Optional token coordinates with shape ``[N, 2]`` or ``[B, N, 2]``. If omitted, a normalized visual grid is inferred. **kwargs: Processor output fields when ``inputs`` is not provided. Returns: Normalized compressed image tokens with shape ``[B, K, D]``. """ model_inputs = dict(inputs) if inputs is not None else kwargs outputs = self.retriever(**model_inputs) embeddings = F.normalize(_extract_embeddings(outputs).float(), p=2, dim=-1) merged = [] for batch_idx, tokens in enumerate(embeddings): item_coords = _select_coords(coords, batch_idx, tokens, self.merge_config) merged.append(merge_tokens(tokens, item_coords, self.merge_config)["tokens"]) return torch.stack(merged, dim=0) @torch.inference_mode() def encode_query(self, inputs=None, return_mask=False, **kwargs): """Encode text queries into normalized multi-vector query tokens. Args: inputs: Optional processor output dictionary. return_mask: Whether to also return a boolean attention mask. **kwargs: Processor output fields when ``inputs`` is not provided. Returns: Query token embeddings, and optionally a valid-token mask. """ model_inputs = dict(inputs) if inputs is not None else kwargs outputs = self.retriever(**model_inputs) embeddings = F.normalize(_extract_embeddings(outputs).float(), p=2, dim=-1) if not return_mask: return embeddings mask = model_inputs.get("attention_mask") if mask is None or mask.shape[:2] != embeddings.shape[:2]: mask = torch.ones(embeddings.shape[:2], dtype=torch.bool, device=embeddings.device) return embeddings, mask.bool() def score(self, query_tokens, image_tokens, query_mask=None, image_mask=None): """Compute ColPali-style mean MaxSim scores.""" sim = torch.einsum("qmd,ind->qimn", query_tokens.float(), image_tokens.float()) if image_mask is not None: sim = sim.masked_fill(~image_mask[None, :, None, :].bool(), torch.finfo(sim.dtype).min) maxsim = sim.max(dim=-1).values if query_mask is None: return maxsim.mean(dim=-1) mask = query_mask.to(dtype=maxsim.dtype, device=maxsim.device) return (maxsim * mask[:, None, :]).sum(dim=-1) / mask.sum(dim=-1).clamp_min(1.0)[:, None] def _extract_embeddings(outputs): for name in ("embeddings", "last_hidden_state"): value = getattr(outputs, name, None) if value is not None: return value if isinstance(outputs, dict): for name in ("embeddings", "last_hidden_state"): if name in outputs: return outputs[name] if isinstance(outputs, (tuple, list)) and outputs: return outputs[0] raise RuntimeError("Could not find token embeddings in model outputs.") def _resolve_weights_path(pretrained_model_name_or_path, **kwargs) -> str: local_path = Path(str(pretrained_model_name_or_path)) / "model.safetensors" if local_path.exists(): return str(local_path) return hf_hub_download( repo_id=str(pretrained_model_name_or_path), filename="model.safetensors", **kwargs, ) def _find_projector(model: nn.Module) -> nn.Module: target = getattr(model, "embedding_proj_layer", None) if target is not None: return target base_model = getattr(model, "base_model", None) target = getattr(base_model, "embedding_proj_layer", None) if base_model is not None else None if target is not None: return target for name, module in model.named_modules(): if "embedding_proj_layer" in name and isinstance(module, nn.Linear): return module raise RuntimeError("Could not find embedding_proj_layer in the base model.") def _select_coords(coords, batch_idx: int, tokens: torch.Tensor, config: MergeConfig) -> torch.Tensor: if coords is None: return make_visual_coords(tokens.size(0), config.image_token_count, tokens.device, tokens.dtype) if coords.dim() == 3: coords = coords[batch_idx] return coords.to(device=tokens.device, dtype=tokens.dtype) def make_visual_coords(num_tokens: int, image_token_count: int, device, dtype) -> torch.Tensor: grid_tokens = min(num_tokens, image_token_count) side = int(grid_tokens**0.5) if side * side != grid_tokens: side = int(image_token_count**0.5) grid_tokens = min(num_tokens, side * side) ys, xs = torch.meshgrid( torch.linspace(0.0, 1.0, side, device=device, dtype=dtype), torch.linspace(0.0, 1.0, side, device=device, dtype=dtype), indexing="ij", ) coords = torch.full((num_tokens, 2), 0.5, device=device, dtype=dtype) coords[:grid_tokens] = torch.stack([xs.flatten(), ys.flatten()], dim=-1)[:grid_tokens] return coords def merge_tokens(tokens: torch.Tensor, coords: torch.Tensor, config: MergeConfig) -> dict[str, torch.Tensor]: tokens = F.normalize(tokens.float(), p=2, dim=-1) assignments, _, feature_centers, spatial_centers = _feature_spatial_clusters(tokens, coords, config) dist = _feature_spatial_distance(tokens, coords, feature_centers, spatial_centers, config.spatial_weight) weights = torch.softmax(-dist / max(float(config.assignment_temperature), 1e-6), dim=-1) denom = weights.sum(dim=0).clamp_min(1e-6) merged_tokens = F.normalize(weights.T @ tokens / denom[:, None], p=2, dim=-1) merged_coords = weights.T @ coords / denom[:, None] return {"tokens": merged_tokens, "coords": merged_coords, "assignments": assignments} def _feature_spatial_clusters(tokens, coords, config): num_tokens = tokens.size(0) k = min(int(config.num_regions), num_tokens) init_idx = _uniform_indices(num_tokens, k, tokens.device) feature_centers = tokens[init_idx].clone() spatial_centers = coords[init_idx].clone() assignments = torch.zeros(num_tokens, dtype=torch.long, device=tokens.device) for _ in range(max(int(config.cluster_iters), 1)): dist = _feature_spatial_distance(tokens, coords, feature_centers, spatial_centers, config.spatial_weight) assignments = dist.argmin(dim=-1) feature_centers, spatial_centers = _update_centers(tokens, coords, assignments, feature_centers, spatial_centers) return assignments, init_idx, feature_centers, spatial_centers def _feature_spatial_distance(tokens, coords, feature_centers, spatial_centers, spatial_weight): dist = 1.0 - tokens @ feature_centers.T if spatial_weight > 0: dist = dist + float(spatial_weight) * (coords[:, None, :] - spatial_centers[None, :, :]).square().sum(dim=-1) return dist def _update_centers(tokens, coords, assignments, old_features, old_spatial): feature_centers = old_features.clone() spatial_centers = old_spatial.clone() for idx in range(old_features.size(0)): mask = assignments == idx if bool(mask.any()): feature_centers[idx] = F.normalize(tokens[mask].mean(dim=0), p=2, dim=-1) spatial_centers[idx] = coords[mask].mean(dim=0) return feature_centers, spatial_centers def _uniform_indices(num_tokens: int, k: int, device: torch.device) -> torch.Tensor: if k == 1: return torch.zeros(1, dtype=torch.long, device=device) return torch.linspace(0, num_tokens - 1, k, device=device).round().long().unique(sorted=True)