| from __future__ import annotations |
|
|
| from typing import Optional |
|
|
| import torch |
| from torch import nn |
| from transformers.models.modernbert.modeling_modernbert import ( |
| ModernBertForMaskedLM, |
| ModernBertModel, |
| ModernBertPredictionHead, |
| ModernBertPreTrainedModel, |
| ) |
|
|
|
|
| class ModernGenaModel(ModernBertModel): |
| """Keep unpadded Flash outputs rank-2 when called by the MLM wrapper.""" |
|
|
| 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.LongTensor] = None, |
| inputs_embeds: 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, |
| ): |
| use_return_dict = ( |
| return_dict if return_dict is not None else self.config.use_return_dict |
| ) |
| output = super().forward( |
| 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=use_return_dict, |
| ) |
|
|
| |
| |
| |
| |
| needs_squeeze = ( |
| self.config._attn_implementation == "flash_attention_2" |
| and indices is not None |
| and output[0].dim() == 3 |
| and output[0].shape[0] == 1 |
| ) |
| if needs_squeeze: |
| if use_return_dict: |
| output.last_hidden_state = output.last_hidden_state.squeeze(0) |
| else: |
| output = (output[0].squeeze(0),) + output[1:] |
| return output |
|
|
|
|
| class ModernGenaForMaskedLM(ModernBertForMaskedLM): |
| """ModernBERT MLM with correct Flash hidden-state/logit repadding.""" |
|
|
| def __init__(self, config): |
| ModernBertPreTrainedModel.__init__(self, config) |
| self.config = config |
| self.model = ModernGenaModel(config) |
| self.head = ModernBertPredictionHead(config) |
| self.decoder = nn.Linear( |
| config.hidden_size, |
| config.vocab_size, |
| bias=config.decoder_bias, |
| ) |
| self.sparse_prediction = config.sparse_prediction |
| self.sparse_pred_ignore_index = config.sparse_pred_ignore_index |
| self.post_init() |
|
|
|
|
| __all__ = ["ModernGenaForMaskedLM", "ModernGenaModel"] |
|
|
