import torch import torch.nn as nn from transformers import ( AutoConfig, AutoModel, AutoModelForMaskedLM, AutoTokenizer, PretrainedConfig, PreTrainedModel, ) from transformers.activations import GELUActivation from transformers.utils import logging logger = logging.get_logger(__name__) class MILCOConfig(PretrainedConfig): model_type = "milco" def __init__(self, **kwargs): super().__init__(**kwargs) self.lsr_encoder_checkpoint = kwargs.get("lsr_encoder_checkpoint", "bert-base-uncased") self.multilingual_encoder_checkpoint = kwargs.get("multilingual_encoder_checkpoint", "bert-base-uncased") self.model_variant = kwargs.get("model_variant", "bert") # Sub-model dimensions — cached in config.json after first save. self.en_vocab_size = kwargs.get("en_vocab_size", None) self.en_hidden_size = kwargs.get("en_hidden_size", None) self.m_vocab_size = kwargs.get("m_vocab_size", None) self.m_hidden_size = kwargs.get("m_hidden_size", None) # Sub-model config dicts — stored so __init__ can rebuild architecture offline. self.en_sub_config = kwargs.get("en_sub_config", None) self.m_sub_config = kwargs.get("m_sub_config", None) def _resolve_sub_configs(self): """Fetch sub-model configs from HF hub if not already cached.""" if self.en_sub_config is None: en_config = AutoConfig.from_pretrained(self.lsr_encoder_checkpoint, trust_remote_code=True) self.en_vocab_size = en_config.vocab_size self.en_hidden_size = en_config.hidden_size self.en_sub_config = en_config.to_dict() if self.m_sub_config is None: m_config = AutoConfig.from_pretrained(self.multilingual_encoder_checkpoint, trust_remote_code=True) self.m_vocab_size = m_config.vocab_size self.m_hidden_size = m_config.hidden_size self.m_sub_config = m_config.to_dict() class MILCOModel(PreTrainedModel): config_class = MILCOConfig base_model_prefix = "milco" def __init__(self, config: MILCOConfig): super().__init__(config) config._resolve_sub_configs() # Build architecture from cached configs — no pretrained weight downloads. m_sub_config = AutoConfig.for_model(**config.m_sub_config) self.m_model = AutoModel.from_config(m_sub_config) en_sub_config = AutoConfig.for_model(**config.en_sub_config) en_lsr = AutoModelForMaskedLM.from_config(en_sub_config) if hasattr(en_lsr, "cls"): self.mlm_head = en_lsr.cls else: self.mlm_head = en_lsr.lm_head self.projector = nn.Linear(config.m_hidden_size, config.en_hidden_size) self.activation = GELUActivation() self.echo = nn.Linear(config.m_hidden_size, 1) self.scale = nn.Parameter(torch.tensor([1.0])) self.post_init() @property def en_tokenizer(self): if not hasattr(self, "_en_tokenizer"): self._en_tokenizer = AutoTokenizer.from_pretrained( self.config.lsr_encoder_checkpoint, trust_remote_code=True, ) return self._en_tokenizer @property def m_tokenizer(self): if not hasattr(self, "_m_tokenizer"): self._m_tokenizer = AutoTokenizer.from_pretrained( self.config.multilingual_encoder_checkpoint, trust_remote_code=True, ) return self._m_tokenizer def get_vocab(self): en_term2id = self.en_tokenizer.vocab en_id2term = {idx: f"e_{term}" for term, idx in en_term2id.items()} m_term2id = self.m_tokenizer.vocab m_id2term = {idx + self.config.en_vocab_size: f"m_{term}" for term, idx in m_term2id.items()} return en_id2term | m_id2term def encode_query(self, *args, **kwargs): return self.encode_text(*args, **kwargs) def encode_document(self, *args, **kwargs): return self.encode_text(*args, **kwargs) def encode_text(self, texts, batch_size: int = 32, max_length: int = None, source_view=False, return_dict=False, **kwargs): if isinstance(texts, str): texts = [texts] if max_length is None: max_length = self.m_tokenizer.model_max_length device = "cuda" if torch.cuda.is_available() else "cpu" self.to(device) reps = [] for idx in range(0, len(texts), batch_size): batch_texts = texts[idx:idx + batch_size] inputs = self.m_tokenizer( batch_texts, padding=True, truncation=True, pad_to_multiple_of=32, max_length=max_length, return_tensors="pt", ) inputs = {k: v.to(device) for k, v in inputs.items()} with torch.no_grad(), torch.autocast(device_type=device, dtype=torch.bfloat16): output = self._encode( input_ids=inputs["input_ids"], attention_mask=inputs["attention_mask"], source_view=source_view, ) if isinstance(output, tuple): sparse_rep = self._build_dual_view_sparse(output) else: sparse_rep = output.to_sparse() reps.append(sparse_rep.cpu()) combined = sparse_row_cat_2d(reps) assert len(texts) == combined.size(0) if return_dict: return self._sparse_to_dicts(combined) return combined def _sparse_to_dicts(self, sparse_tensor): """Convert a sparse [N, vocab] tensor to a list of {token: weight} dicts.""" id2term = self.get_vocab() sparse_tensor = sparse_tensor.coalesce() indices = sparse_tensor.indices() values = sparse_tensor.values() rows, cols = indices[0], indices[1] result = [[] for _ in range(sparse_tensor.size(0))] for i in range(len(values)): row = rows[i].item() col = cols[i].item() weight = values[i].item() token = id2term.get(col, f"unk_{col}") result[row].append((token, weight)) return [ dict(sorted(pairs, key=lambda x: x[1], reverse=True)) for pairs in result ] def _build_dual_view_sparse(self, output): pivot_view, (source_ids, source_weights) = output batch_size = source_ids.size(0) vocab_size = self.config.en_vocab_size + self.config.m_vocab_size pivot_rows, pivot_cols = pivot_view.nonzero(as_tuple=True) pivot_vals = pivot_view[pivot_rows, pivot_cols].contiguous() source_rows = torch.arange(batch_size, device=source_ids.device).repeat_interleave(source_ids.size(1)) source_cols = source_ids.view(-1) source_vals = source_weights.view(-1) nonzero_mask = source_vals > 0 source_rows = source_rows[nonzero_mask] source_cols = source_cols[nonzero_mask] + self.config.en_vocab_size source_vals = source_vals[nonzero_mask] rows = torch.cat([pivot_rows, source_rows]) cols = torch.cat([pivot_cols, source_cols]) indices = torch.stack([rows, cols], dim=0) values = torch.cat([pivot_vals, source_vals]) return torch.sparse_coo_tensor(indices, values, (batch_size, vocab_size)) def _encode(self, input_ids, attention_mask, source_view=False): hidden_states_m = self.m_model(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state hidden_states_e = self.activation(self.projector(hidden_states_m)) logits = self.mlm_head(hidden_states_e) * attention_mask.unsqueeze(-1) pivot_view = torch.log1p(torch.relu(logits.max(dim=1).values)) if source_view: token_scores = torch.relu(self.echo(hidden_states_m).squeeze(-1)) source_weights = token_scores * attention_mask return pivot_view, (input_ids, source_weights) return pivot_view def forward(self, *args, **kwargs): return self._encode(*args, **kwargs) def save_pretrained(self, save_directory: str, **kwargs): kwargs["safe_serialization"] = False super().save_pretrained(save_directory, **kwargs) def sparse_row_cat_2d(tensors): """Concatenate a list of 2-D COO sparse tensors along rows (dim=0).""" assert len(tensors) > 0, "empty list" if len(tensors) == 1: return tensors[0] tensors = [t.coalesce() for t in tensors] assert all(t.layout == torch.sparse_coo for t in tensors), "COO only" ncols = tensors[0].size(1) assert all(t.dim() == 2 and t.size(1) == ncols for t in tensors), "size(1) mismatch" total_rows = sum(t.size(0) for t in tensors) device = tensors[0].device ind_parts, val_parts = [], [] row_offset = 0 for t in tensors: idx = t._indices() if row_offset: idx = idx.clone() idx[0] += row_offset ind_parts.append(idx) val_parts.append(t._values()) row_offset += t.size(0) indices = torch.cat(ind_parts, dim=1) values = torch.cat(val_parts, dim=0) return torch.sparse_coo_tensor(indices, values, (total_rows, ncols), device=device).coalesce()