""" ColVec1 - ColVec1 retrieval wrapper for late interaction. """ import glob import json import os from typing import ClassVar, List, Optional import torch import torch.nn as nn from transformers import AutoModelForImageTextToText, PreTrainedModel from .configuration_colvec1 import ColVec1Config class ColVec1PreTrainedModel(PreTrainedModel): """Base class for ColVec1 models.""" config_class = ColVec1Config base_model_prefix = "colvec1" supports_gradient_checkpointing = True _tied_weights_keys: ClassVar[List[str]] = [] class ColVec1(ColVec1PreTrainedModel): """ Retrieval model wrapper for ColVec1 checkpoints. It loads the upstream model with `AutoModelForImageTextToText`, then adds a projection head to produce L2-normalized retrieval embeddings. """ main_input_name: ClassVar[str] = "input_ids" def __init__(self, config: ColVec1Config): super().__init__(config) self.config = config self.vlm = None self.embedding_proj_layer = nn.Linear(config.text_hidden_size, config.embed_dim) self.post_init() def forward( self, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, pixel_values: Optional[torch.FloatTensor] = None, **kwargs, ) -> torch.Tensor: kwargs.pop("output_hidden_states", None) kwargs.pop("return_dict", None) vlm_outputs = self.vlm( input_ids=input_ids, attention_mask=attention_mask, pixel_values=pixel_values, output_hidden_states=True, return_dict=True, **kwargs, ) if hasattr(vlm_outputs, "hidden_states") and vlm_outputs.hidden_states is not None: last_hidden_states = vlm_outputs.hidden_states[-1] elif hasattr(vlm_outputs, "last_hidden_state"): last_hidden_states = vlm_outputs.last_hidden_state else: last_hidden_states = vlm_outputs[0] embeddings = self.embedding_proj_layer( last_hidden_states.to(self.embedding_proj_layer.weight.dtype) ) embeddings = nn.functional.normalize(embeddings, p=2, dim=-1) if attention_mask is not None: embeddings = embeddings * attention_mask.unsqueeze(-1) return embeddings @classmethod def from_pretrained( cls, pretrained_model_name_or_path: str, embed_dim: int = 128, torch_dtype: torch.dtype = None, device_map: str = None, attn_impl: str = None, **kwargs, ): # AutoModel may rename torch_dtype -> dtype in newer transformers if torch_dtype is None: torch_dtype = kwargs.pop("dtype", None) # Pop config early so we can inspect model_type for merged-repo detection. # When called via AutoModel.from_pretrained, transformers resolves the config # and passes it here as a kwarg; config = kwargs.pop("config", None) if config is not None and hasattr(config, "embed_dim"): embed_dim = config.embed_dim # Detect a merged ColVec1 repo using three strategies in order: # 1. config object already provided (Hub path via AutoModel dispatch) # 2. local config.json on disk (direct local-path usage) # 3. AutoConfig.from_pretrained (direct Hub ID usage without AutoModel) _is_merged = ( config is not None and getattr(config, "model_type", None) == "colvec1" ) if not _is_merged: config_path = os.path.join(pretrained_model_name_or_path, "config.json") if os.path.exists(config_path): with open(config_path) as f: raw = json.load(f) _is_merged = raw.get("model_type") == "colvec1" else: # Remote Hub ID: fetch the config to check model_type. from transformers import AutoConfig try: hub_config = AutoConfig.from_pretrained( pretrained_model_name_or_path, trust_remote_code=kwargs.get("trust_remote_code", True), ) _is_merged = getattr(hub_config, "model_type", None) == "colvec1" except Exception: pass if _is_merged: return cls._load_merged( pretrained_model_name_or_path, torch_dtype=torch_dtype, device_map=device_map, attn_impl=attn_impl, **kwargs, ) # --- From-scratch path: load a raw Qwen3.5 VLM and wrap it --- # (config was already popped above; rest of the method is unchanged) vlm_kwargs = {"trust_remote_code": kwargs.pop("trust_remote_code", True)} if torch_dtype is not None: vlm_kwargs["torch_dtype"] = torch_dtype if device_map is not None: vlm_kwargs["device_map"] = device_map if attn_impl is not None: vlm_kwargs["attn_implementation"] = attn_impl if "quantization_config" in kwargs: vlm_kwargs["quantization_config"] = kwargs.pop("quantization_config") vlm = AutoModelForImageTextToText.from_pretrained(pretrained_model_name_or_path, **vlm_kwargs) if hasattr(vlm.config, "text_config") and hasattr(vlm.config.text_config, "hidden_size"): text_hidden_size = vlm.config.text_config.hidden_size else: text_hidden_size = getattr(vlm.config, "hidden_size", 2560) model_config = ColVec1Config( embed_dim=embed_dim, text_hidden_size=text_hidden_size, padding_side="left", ) model = cls(model_config) model.vlm = vlm model.embedding_proj_layer = nn.Linear(model_config.text_hidden_size, model_config.embed_dim) if torch_dtype is not None: model.embedding_proj_layer = model.embedding_proj_layer.to(torch_dtype) if hasattr(vlm, "device"): model.embedding_proj_layer = model.embedding_proj_layer.to(vlm.device) tied = getattr(vlm, "_tied_weights_keys", None) if isinstance(tied, dict): model._tied_weights_keys = {f"vlm.{k}": f"vlm.{v}" for k, v in tied.items()} elif isinstance(tied, (list, tuple, set)): model._tied_weights_keys = [f"vlm.{k}" for k in tied] else: model._tied_weights_keys = [] return model @classmethod def _load_merged( cls, path: str, torch_dtype: torch.dtype = None, device_map: str = None, attn_impl: str = None, **kwargs, ): """Load a merged ColVec1 checkpoint (dense VLM weights + embedding_proj_layer).""" from safetensors.torch import load_file # Resolve Hub repo ID to a local cached snapshot directory so all # subsequent os.path / glob operations work for both local and remote paths. if not os.path.isdir(path): from huggingface_hub import snapshot_download path = snapshot_download(path) config = ColVec1Config.from_pretrained(path) base_name = config.base_model_name_or_path if base_name is None: raise ValueError( f"Merged ColVec1 config at {path} is missing 'base_model_name_or_path'. " "This field is required to know which VLM architecture to instantiate." ) vlm_kwargs = {"trust_remote_code": True} if torch_dtype is not None: vlm_kwargs["torch_dtype"] = torch_dtype if device_map is not None: vlm_kwargs["device_map"] = device_map if attn_impl is not None: vlm_kwargs["attn_implementation"] = attn_impl vlm = AutoModelForImageTextToText.from_pretrained(base_name, **vlm_kwargs) model = cls(config) model.vlm = vlm safetensor_files = sorted(glob.glob(os.path.join(path, "model*.safetensors"))) if not safetensor_files: raise FileNotFoundError(f"No model*.safetensors files found in {path}") state_dict = {} for sf in safetensor_files: state_dict.update(load_file(sf)) model.load_state_dict(state_dict, strict=False) if torch_dtype is not None: model.embedding_proj_layer = model.embedding_proj_layer.to(torch_dtype) if hasattr(vlm, "device"): model.embedding_proj_layer = model.embedding_proj_layer.to(vlm.device) tied = getattr(vlm, "_tied_weights_keys", None) if isinstance(tied, dict): model._tied_weights_keys = {f"vlm.{k}": f"vlm.{v}" for k, v in tied.items()} elif isinstance(tied, (list, tuple, set)): model._tied_weights_keys = [f"vlm.{k}" for k in tied] else: model._tied_weights_keys = [] return model def tie_weights(self, *args, **kwargs): if self.vlm is None: # Called during post_init() before the wrapped VLM is attached. return None try: return self.vlm.tie_weights(*args, **kwargs) except TypeError: return self.vlm.tie_weights() def get_input_embeddings(self): return self.vlm.get_input_embeddings() def set_input_embeddings(self, value): self.vlm.set_input_embeddings(value) def get_output_embeddings(self): return self.vlm.get_output_embeddings() def set_output_embeddings(self, new_embeddings): self.vlm.set_output_embeddings(new_embeddings) def resize_token_embeddings( self, new_num_tokens: Optional[int] = None, pad_to_multiple_of: Optional[int] = None, mean_resizing: bool = True, ) -> nn.Embedding: model_embeds = self.vlm.resize_token_embeddings( new_num_tokens=new_num_tokens, pad_to_multiple_of=pad_to_multiple_of, mean_resizing=mean_resizing, ) if hasattr(self.vlm.config, "text_config"): self.vlm.config.text_config.vocab_size = model_embeds.num_embeddings if hasattr(self.vlm.config, "vocab_size"): self.vlm.config.vocab_size = model_embeds.num_embeddings return model_embeds @property def device(self): return next(self.parameters()).device def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None): if self.vlm is not None and hasattr(self.vlm, "gradient_checkpointing_enable"): self.vlm.gradient_checkpointing_enable(gradient_checkpointing_kwargs) def gradient_checkpointing_disable(self): if self.vlm is not None and hasattr(self.vlm, "gradient_checkpointing_disable"): self.vlm.gradient_checkpointing_disable() __all__ = ["ColVec1", "ColVec1PreTrainedModel"]