| import torch |
| import transformers |
|
|
| from typing import Optional, Union |
| from transformers.modeling_outputs import MaskedLMOutput |
|
|
| class ModelForMaskedLM(transformers.ModernBertForMaskedLM): |
|
|
| def __init__(self, config: transformers.ModernBertConfig): |
| super().__init__(config) |
| self.decoder = torch.nn.Linear(config.hidden_size, config.num_concepts, bias=config.decoder_bias) |
|
|
| |
| self.post_init() |
|
|
| if hasattr(self.model.config, 'post_cls'): |
| self.post_cls = torch.nn.Linear(config.num_concepts, config.vocab_size, bias=config.decoder_bias) |
|
|
| def forward( |
| self, |
| input_ids: Optional[torch.LongTensor] = None, |
| attention_mask: Optional[torch.Tensor] = None, |
| sliding_window_mask: Optional[torch.Tensor] = None, |
| position_ids: Optional[torch.Tensor] = None, |
| inputs_embeds: Optional[torch.Tensor] = None, |
| labels: Optional[torch.Tensor] = None, |
| indices: Optional[torch.Tensor] = None, |
| cu_seqlens: Optional[torch.Tensor] = None, |
| max_seqlen: Optional[int] = None, |
| batch_size: Optional[int] = None, |
| seq_len: Optional[int] = None, |
| output_attentions: Optional[bool] = None, |
| output_hidden_states: Optional[bool] = None, |
| return_dict: Optional[bool] = None, |
| **kwargs, |
| ) -> Union[tuple[torch.Tensor], MaskedLMOutput]: |
| r""" |
| sliding_window_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): |
| Mask to avoid performing attention on padding or far-away tokens. In ModernBert, only every few layers |
| perform global attention, while the rest perform local attention. This mask is used to avoid attending to |
| far-away tokens in the local attention layers when not using Flash Attention. |
| indices (`torch.Tensor` of shape `(total_unpadded_tokens,)`, *optional*): |
| Indices of the non-padding tokens in the input sequence. Used for unpadding the output. |
| cu_seqlens (`torch.Tensor` of shape `(batch + 1,)`, *optional*): |
| Cumulative sequence lengths of the input sequences. Used to index the unpadded tensors. |
| max_seqlen (`int`, *optional*): |
| Maximum sequence length in the batch excluding padding tokens. Used to unpad input_ids and pad output tensors. |
| batch_size (`int`, *optional*): |
| Batch size of the input sequences. Used to pad the output tensors. |
| seq_len (`int`, *optional*): |
| Sequence length of the input sequences including padding tokens. Used to pad the output tensors. |
| """ |
| return_dict = return_dict if return_dict is not None else self.config.use_return_dict |
| self._maybe_set_compile() |
|
|
| if self.config._attn_implementation == "flash_attention_2": |
| if indices is None and cu_seqlens is None and max_seqlen is None: |
| if batch_size is None and seq_len is None: |
| if inputs_embeds is not None: |
| batch_size, seq_len = inputs_embeds.shape[:2] |
| else: |
| batch_size, seq_len = input_ids.shape[:2] |
| device = input_ids.device if input_ids is not None else inputs_embeds.device |
|
|
| if attention_mask is None: |
| attention_mask = torch.ones((batch_size, seq_len), device=device, dtype=torch.bool) |
|
|
| if inputs_embeds is None: |
| with torch.no_grad(): |
| input_ids, indices, cu_seqlens, max_seqlen, position_ids, labels = _unpad_modernbert_input( |
| inputs=input_ids, attention_mask=attention_mask, position_ids=position_ids, labels=labels |
| ) |
| else: |
| inputs_embeds, indices, cu_seqlens, max_seqlen, position_ids, labels = _unpad_modernbert_input( |
| inputs=inputs_embeds, attention_mask=attention_mask, position_ids=position_ids, labels=labels |
| ) |
|
|
| outputs = self.model( |
| input_ids=input_ids, |
| attention_mask=attention_mask, |
| sliding_window_mask=sliding_window_mask, |
| position_ids=position_ids, |
| inputs_embeds=inputs_embeds, |
| indices=indices, |
| cu_seqlens=cu_seqlens, |
| max_seqlen=max_seqlen, |
| batch_size=batch_size, |
| seq_len=seq_len, |
| output_attentions=output_attentions, |
| output_hidden_states=output_hidden_states, |
| return_dict=return_dict, |
| ) |
| last_hidden_state = outputs[0] |
|
|
| if self.sparse_prediction and labels is not None: |
| |
| labels = labels.view(-1) |
| last_hidden_state = last_hidden_state.view(labels.shape[0], -1) |
|
|
| |
| mask_tokens = labels != self.sparse_pred_ignore_index |
| last_hidden_state = last_hidden_state[mask_tokens] |
| labels = labels[mask_tokens] |
|
|
| logits = ( |
| self.compiled_head(last_hidden_state) |
| if self.config.reference_compile |
| else self.decoder(self.head(last_hidden_state)) |
| ) |
|
|
| if hasattr(self, 'post_cls'): |
| logits = self.post_cls(logits) |
|
|
| loss = None |
| if labels is not None: |
| loss = self.loss_function(logits, labels, vocab_size=self.config.vocab_size, **kwargs) |
|
|
| if self.config._attn_implementation == "flash_attention_2": |
| |
| with nullcontext() if self.config.repad_logits_with_grad or labels is None else torch.no_grad(): |
| logits = _pad_modernbert_output(inputs=logits, indices=indices, batch=batch_size, seqlen=seq_len) |
| |
| if getattr(outputs, "hidden_states", None) is not None: |
| padded_hidden_states = [] |
| for hs in outputs.hidden_states: |
| if hs.dim() == 3 and hs.shape[0] == 1: |
| hs = hs.squeeze(0) |
| padded_hidden_states.append( |
| _pad_modernbert_output(inputs=hs, indices=indices, batch=batch_size, seqlen=seq_len) |
| ) |
| outputs.hidden_states = tuple(padded_hidden_states) |
|
|
| if not return_dict: |
| output = (logits,) |
| return ((loss,) + output) if loss is not None else output |
|
|
| return MaskedLMOutput( |
| loss=loss, |
| logits=logits, |
| hidden_states=outputs.hidden_states, |
| attentions=outputs.attentions, |
| ) |
|
|
