"""Self-contained model definition for the MARS shared-cross-attention checkpoint. This is a faithful reproduction of `SharedCrossAttentionModel` from the `summ-mask-benchmark` research repo, stripped of optional research flags (copy mechanism, type bias) that are NOT used by this checkpoint. Architecture: a single ModernBERT-base encoder shared between the summary and the masked text, followed by 2 cross-attention layers where the masked text attends to the summary, then a vocab projection. """ import math from typing import Optional import torch import torch.nn as nn import torch.nn.functional as F from transformers import AutoModel from transformers.modeling_outputs import MaskedLMOutput ENTMASK = "[ENTMASK]" ENTSTART = "[ENTSTART]" ENTEND = "[ENTEND]" ENTITY_TYPES = [ "PERSON", "ORG", "GPE", "LOC", "DATE", "TIME", "MONEY", "QUANTITY", "PERCENT", "CARDINAL", "ORDINAL", "EVENT", "WORK_OF_ART", "LAW", "LANGUAGE", "FAC", "PRODUCT", "NORP", ] TYPED_ENTMASK_TOKENS = {t: f"[ENTMASK_{t}]" for t in ENTITY_TYPES} class CrossAttentionLayer(nn.Module): def __init__(self, hidden_size: int, num_attention_heads: int = 12, attention_dropout: float = 0.1): super().__init__() self.num_attention_heads = num_attention_heads self.attention_head_size = hidden_size // num_attention_heads self.all_head_size = self.num_attention_heads * self.attention_head_size self.query = nn.Linear(hidden_size, self.all_head_size) self.key = nn.Linear(hidden_size, self.all_head_size) self.value = nn.Linear(hidden_size, self.all_head_size) self.dropout = nn.Dropout(attention_dropout) self.dense = nn.Linear(hidden_size, hidden_size) self.layer_norm = nn.LayerNorm(hidden_size) def _shape(self, x: torch.Tensor) -> torch.Tensor: new_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size) return x.view(*new_shape).permute(0, 2, 1, 3) def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, encoder_attention_mask: Optional[torch.Tensor] = None, ) -> torch.Tensor: q = self._shape(self.query(hidden_states)) k = self._shape(self.key(encoder_hidden_states)) v = self._shape(self.value(encoder_hidden_states)) scores = torch.matmul(q, k.transpose(-1, -2)) / math.sqrt(self.attention_head_size) if encoder_attention_mask is not None: ext = encoder_attention_mask[:, None, None, :] ext = (1.0 - ext) * torch.finfo(scores.dtype).min scores = scores + ext probs = self.dropout(F.softmax(scores, dim=-1)) ctx = torch.matmul(probs, v).permute(0, 2, 1, 3).contiguous() ctx = ctx.view(*ctx.size()[:-2], self.all_head_size) out = self.dropout(self.dense(ctx)) return self.layer_norm(hidden_states + out) class SharedCrossAttentionModel(nn.Module): """Shared-encoder cross-attention model for masked entity infilling. Use `summary_input_ids` / `masked_input_ids` (both encoded with the same tokenizer that includes [ENTMASK], [ENTSTART], [ENTEND] and the typed [ENTMASK_] specials). """ def __init__( self, model_name: str = "answerdotai/ModernBERT-base", num_cross_attention_layers: int = 2, num_attention_heads: int = 12, dropout: float = 0.1, ): super().__init__() self.encoder = AutoModel.from_pretrained(model_name) hidden_size = self.encoder.config.hidden_size vocab_size = self.encoder.config.vocab_size encoder_heads = getattr(self.encoder.config, "num_attention_heads", num_attention_heads) if hidden_size % num_attention_heads != 0: num_attention_heads = encoder_heads self.cross_attention_layers = nn.ModuleList( [CrossAttentionLayer(hidden_size, num_attention_heads, dropout) for _ in range(num_cross_attention_layers)] ) self.predictions = nn.Linear(hidden_size, vocab_size) self.hidden_size = hidden_size self.vocab_size = vocab_size def resize_token_embeddings(self, new_num_tokens: int): self.encoder.resize_token_embeddings(new_num_tokens) old = self.predictions self.predictions = nn.Linear(self.hidden_size, new_num_tokens) n = min(old.out_features, new_num_tokens) self.predictions.weight.data[:n] = old.weight.data[:n] self.predictions.bias.data[:n] = old.bias.data[:n] self.vocab_size = new_num_tokens def forward( self, summary_input_ids: torch.Tensor, summary_attention_mask: torch.Tensor, masked_input_ids: torch.Tensor, masked_attention_mask: torch.Tensor, labels: Optional[torch.Tensor] = None, ) -> MaskedLMOutput: s_hid = self.encoder(input_ids=summary_input_ids, attention_mask=summary_attention_mask).last_hidden_state m_hid = self.encoder(input_ids=masked_input_ids, attention_mask=masked_attention_mask).last_hidden_state for cross in self.cross_attention_layers: m_hid = cross(m_hid, s_hid, summary_attention_mask) logits = self.predictions(m_hid) loss = None if labels is not None: loss = nn.CrossEntropyLoss(ignore_index=-100)( logits.view(-1, self.vocab_size), labels.view(-1) ) return MaskedLMOutput(loss=loss, logits=logits, hidden_states=m_hid) def load_model_from_checkpoint( repo_or_path: str, base_model: str = "answerdotai/ModernBERT-base", device: Optional[str] = None, ): """Load the MARS checkpoint and matching tokenizer from a local dir or a Hugging Face repo id. Returns: (model, tokenizer, device_str) """ from pathlib import Path from transformers import AutoTokenizer from huggingface_hub import snapshot_download if device is None: device = "cuda" if torch.cuda.is_available() else ( "mps" if torch.backends.mps.is_available() else "cpu" ) path = Path(repo_or_path) if not path.exists(): path = Path(snapshot_download(repo_id=repo_or_path)) tokenizer = AutoTokenizer.from_pretrained(path, use_fast=True) model = SharedCrossAttentionModel(model_name=base_model) model.resize_token_embeddings(len(tokenizer)) state = torch.load(path / "model.pt", map_location=device, weights_only=True) model.load_state_dict(state) model.to(device).eval() return model, tokenizer, device