Add modeling_mars.py

modeling_mars.py

@@ -0,0 +1,171 @@
+"""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_<TYPE>] 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