"""ColBERT-Ko: transformers-only inference for dragonkue/colbert-ko-0.1b. Usage: from transformers import AutoModel, AutoTokenizer tokenizer = AutoTokenizer.from_pretrained("dragonkue/colbert-ko-0.1b") model = AutoModel.from_pretrained( "dragonkue/colbert-ko-0.1b", trust_remote_code=True, ) q_embs = model.encode(tokenizer, ["검색 쿼리"], is_query=True) d_embs = model.encode(tokenizer, ["문서 내용"], is_query=False) score = model.maxsim(q_embs[0], d_embs[0]) The model is a ModernBert encoder followed by a 768→128 linear projection (no bias) producing L2-normalized token-level embeddings. Late-interaction MaxSim is the similarity function. Query expansion (pad to 32 tokens with MASK) and document skiplist (punctuation removal) are applied as in the original ColBERT recipe. """ from __future__ import annotations import os from typing import Any import torch import torch.nn.functional as F from torch import nn from transformers import PreTrainedModel from transformers.modeling_outputs import BaseModelOutput from transformers.models.modernbert.configuration_modernbert import ModernBertConfig from transformers.models.modernbert.modeling_modernbert import ModernBertModel DEFAULT_SKIPLIST_WORDS = list("!\"#$%&'()*+,-./:;<=>?@[\\]^_`{|}~") class ColBERTKoConfig(ModernBertConfig): """ModernBert config + ColBERT-specific fields.""" model_type = "modernbert" # keep ModernBert mapping for the encoder def __init__( self, embedding_dim: int = 128, query_prefix: str = "[Q] ", document_prefix: str = "[D] ", query_length: int = 32, document_length: int = 2048, attend_to_expansion_tokens: bool = False, do_query_expansion: bool = True, skiplist_words: list[str] | None = None, **kwargs: Any, ) -> None: super().__init__(**kwargs) self.embedding_dim = embedding_dim self.query_prefix = query_prefix self.document_prefix = document_prefix self.query_length = query_length self.document_length = document_length self.attend_to_expansion_tokens = attend_to_expansion_tokens self.do_query_expansion = do_query_expansion self.skiplist_words = ( skiplist_words if skiplist_words is not None else DEFAULT_SKIPLIST_WORDS ) class ColBERTKo(PreTrainedModel): """ModernBert encoder + 768→128 Linear (no bias) + L2 norm.""" config_class = ColBERTKoConfig base_model_prefix = "encoder" _no_split_modules = ["ModernBertEncoderLayer"] # Linear projection weights are loaded separately from matryoshka_heads.pt # (or 1_Dense/model.safetensors as a fallback); suppress the spurious # "missing key" warning from transformers' standard loader. _keys_to_ignore_on_load_missing = ["linear.weight"] def __init__(self, config: ColBERTKoConfig) -> None: super().__init__(config) self.encoder = ModernBertModel(config) self.linear = nn.Linear( config.hidden_size, config.embedding_dim, bias=False ) # Matryoshka heads (32/64/96/128) populated lazily from matryoshka_heads.pt. # Active dim defaults to embedding_dim (128). set_active_dim() switches # which weight is loaded into self.linear. self._matryoshka_heads: dict[int, torch.Tensor] = {} self._active_dim: int = config.embedding_dim # token id caches resolved at first encode() call self._query_prefix_id: int | None = None self._document_prefix_id: int | None = None self._skiplist_ids: list[int] = [] # Required by transformers >= 5.x for tied-weight bookkeeping. self.post_init() # ---- weight loading (linear comes from matryoshka_heads.pt) ---- @classmethod def from_pretrained( # type: ignore[override] cls, pretrained_model_name_or_path: str | os.PathLike, *args: Any, **kwargs: Any, ) -> "ColBERTKo": model: "ColBERTKo" = super().from_pretrained( # type: ignore[assignment] pretrained_model_name_or_path, *args, **kwargs ) model._load_matryoshka_heads(pretrained_model_name_or_path, **kwargs) model.set_active_dim(model.config.embedding_dim) return model def _load_matryoshka_heads( self, path: str | os.PathLike, **kwargs: Any ) -> None: """Load all Matryoshka projection heads (32/64/96/128) from matryoshka_heads.pt. Falls back to 1_Dense/model.safetensors (128-dim only) if matryoshka_heads.pt is unavailable. """ local_mat = os.path.join(str(path), "matryoshka_heads.pt") local_1d = os.path.join(str(path), "1_Dense", "model.safetensors") heads_path: str | None = None dense_path: str | None = None if os.path.isdir(str(path)): if os.path.exists(local_mat): heads_path = local_mat if os.path.exists(local_1d): dense_path = local_1d else: from huggingface_hub import hf_hub_download hf_kwargs: dict[str, Any] = {} for k in ("revision", "cache_dir", "token", "subfolder"): if k in kwargs: hf_kwargs[k] = kwargs[k] try: heads_path = hf_hub_download( repo_id=str(path), filename="matryoshka_heads.pt", **hf_kwargs ) except Exception: heads_path = None try: dense_path = hf_hub_download( repo_id=str(path), filename="1_Dense/model.safetensors", **hf_kwargs, ) except Exception: dense_path = None if heads_path is not None: ckpt = torch.load(heads_path, map_location="cpu", weights_only=False) for k, v in ckpt["projection_heads"].items(): dim = int(str(k).split(".")[0]) self._matryoshka_heads[dim] = v.to(self.linear.weight.dtype) if not self._matryoshka_heads and dense_path is not None: from safetensors.torch import load_file sd = load_file(dense_path) self._matryoshka_heads[ self.config.embedding_dim ] = sd["linear.weight"].to(self.linear.weight.dtype) if not self._matryoshka_heads: raise FileNotFoundError( f"Could not load projection weights from {path!r}. " f"Expected matryoshka_heads.pt or 1_Dense/model.safetensors." ) def set_active_dim(self, dim: int) -> None: """Switch the active Matryoshka embedding dimension.""" if dim not in self._matryoshka_heads: available = sorted(self._matryoshka_heads.keys()) raise ValueError( f"dim={dim} not available. Available dims: {available}" ) w = self._matryoshka_heads[dim] with torch.no_grad(): new_linear = nn.Linear(w.size(1), w.size(0), bias=False).to( device=self.linear.weight.device, dtype=self.linear.weight.dtype ) new_linear.weight.copy_(w) self.linear = new_linear self._active_dim = dim self.config.embedding_dim = dim @property def active_dim(self) -> int: return self._active_dim @property def matryoshka_dims(self) -> list[int]: return sorted(self._matryoshka_heads.keys()) # ---- tokenization helpers ---- def _resolve_special_ids(self, tokenizer: Any) -> None: if self._query_prefix_id is None: q = self.config.query_prefix for cand in (q, q.strip(), q.rstrip(), q.replace(" ", "")): tid = tokenizer.convert_tokens_to_ids(cand) if tid is not None and tid != tokenizer.unk_token_id: self._query_prefix_id = int(tid) break if self._document_prefix_id is None: d = self.config.document_prefix for cand in (d, d.strip(), d.rstrip(), d.replace(" ", "")): tid = tokenizer.convert_tokens_to_ids(cand) if tid is not None and tid != tokenizer.unk_token_id: self._document_prefix_id = int(tid) break if not self._skiplist_ids: self._skiplist_ids = [ tokenizer.convert_tokens_to_ids(w) for w in self.config.skiplist_words ] def _tokenize( self, tokenizer: Any, texts: list[str], is_query: bool ) -> dict[str, torch.Tensor]: self._resolve_special_ids(tokenizer) max_length = ( self.config.query_length if is_query else self.config.document_length ) prefix_id = self._query_prefix_id if is_query else self._document_prefix_id assert prefix_id is not None, "Failed to resolve ColBERT prefix token id" do_expand = is_query and self.config.do_query_expansion # Queries: pad to max_length for ColBERT mask-token expansion. # Documents: pad to the longest item in the batch so the batch tensor is rectangular. padding_strategy = "max_length" if do_expand else "longest" enc = tokenizer( texts, padding=padding_strategy, truncation=True, max_length=max_length - 1, # reserve 1 slot for prefix token return_tensors="pt", ) input_ids = enc["input_ids"] attn = enc["attention_mask"] # Insert prefix at position 1 (right after BOS/CLS). b = input_ids.size(0) prefix_col = torch.full( (b, 1), prefix_id, dtype=input_ids.dtype ) attn_col = torch.ones((b, 1), dtype=attn.dtype) input_ids = torch.cat([input_ids[:, :1], prefix_col, input_ids[:, 1:]], dim=1) attn = torch.cat([attn[:, :1], attn_col, attn[:, 1:]], dim=1) if is_query and do_expand and self.config.attend_to_expansion_tokens: attn = torch.ones_like(attn) return {"input_ids": input_ids, "attention_mask": attn} # ---- forward ---- def forward( self, input_ids: torch.Tensor, attention_mask: torch.Tensor | None = None, **_: Any, ) -> BaseModelOutput: out = self.encoder(input_ids=input_ids, attention_mask=attention_mask) hidden = out.last_hidden_state if hasattr(out, "last_hidden_state") else out[0] projected = self.linear(hidden) normalized = F.normalize(projected, p=2, dim=-1) return BaseModelOutput(last_hidden_state=normalized) # ---- public API ---- @torch.no_grad() def encode( self, tokenizer: Any, texts: list[str], is_query: bool, batch_size: int = 16, ) -> list[torch.Tensor]: """Return per-text token embeddings (L2-normalized, with skiplist applied).""" self.eval() device = next(self.parameters()).device outputs: list[torch.Tensor] = [] for start in range(0, len(texts), batch_size): batch = texts[start : start + batch_size] feats = self._tokenize(tokenizer, batch, is_query) feats = {k: v.to(device) for k, v in feats.items()} out = self( input_ids=feats["input_ids"], attention_mask=feats["attention_mask"] ) token_embs = out.last_hidden_state # (B, L, d) for i in range(token_embs.size(0)): ids = feats["input_ids"][i] mask = feats["attention_mask"][i].bool() if is_query: keep = ( torch.ones_like(ids, dtype=torch.bool) if self.config.do_query_expansion else mask ) else: keep = mask.clone() skiplist = torch.tensor( self._skiplist_ids, device=ids.device, dtype=ids.dtype, ) keep &= ~torch.isin(ids, skiplist) outputs.append(token_embs[i][keep].cpu()) return outputs @staticmethod def maxsim(q_emb: torch.Tensor, d_emb: torch.Tensor) -> float: """ColBERT MaxSim: Σ_q max_d (q · d). Inputs assumed L2-normalized.""" q = q_emb.to(torch.float32) d = d_emb.to(torch.float32) return float((q @ d.T).max(dim=1).values.sum().item())