File size: 4,800 Bytes
62a610e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 | """Exact Hugging Face-compatible implementation of BashkirRoBERTa Pre-LN."""
import torch
from torch import nn
from transformers import PreTrainedModel
from transformers.modeling_outputs import MaskedLMOutput
try:
from .configuration_bashkir_roberta import BashkirRobertaConfig
except (ImportError, ValueError):
from configuration_bashkir_roberta import BashkirRobertaConfig
class BashkirRobertaEmbeddings(nn.Module):
def __init__(self, config):
super().__init__()
self.word_embeddings = nn.Embedding(
config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id
)
self.position_embeddings = nn.Embedding(
config.max_position_embeddings, config.hidden_size
)
self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=1e-12)
self.dropout = nn.Dropout(config.hidden_dropout_prob)
def forward(self, input_ids):
sequence_length = input_ids.shape[1]
positions = torch.arange(
sequence_length, dtype=torch.long, device=input_ids.device
).unsqueeze(0)
hidden_states = self.word_embeddings(input_ids) + self.position_embeddings(positions)
return self.dropout(self.LayerNorm(hidden_states))
class BashkirRobertaLayer(nn.Module):
"""Pre-LayerNorm transformer block, matching ``train_flagship.py`` exactly."""
def __init__(self, config):
super().__init__()
self.self_attn = nn.MultiheadAttention(
config.hidden_size,
config.num_attention_heads,
dropout=config.hidden_dropout_prob,
batch_first=True,
)
self.norm1 = nn.LayerNorm(config.hidden_size, eps=1e-12)
self.ffn = nn.Sequential(
nn.Linear(config.hidden_size, config.intermediate_size),
nn.GELU(),
nn.Dropout(config.hidden_dropout_prob),
nn.Linear(config.intermediate_size, config.hidden_size),
nn.Dropout(config.hidden_dropout_prob),
)
self.norm2 = nn.LayerNorm(config.hidden_size, eps=1e-12)
def forward(self, hidden_states, padding_mask=None):
normalized = self.norm1(hidden_states)
attention, _ = self.self_attn(
normalized, normalized, normalized, key_padding_mask=padding_mask
)
hidden_states = hidden_states + attention
return hidden_states + self.ffn(self.norm2(hidden_states))
class BashkirRobertaForMaskedLM(PreTrainedModel):
config_class = BashkirRobertaConfig
base_model_prefix = "bashkir_roberta"
main_input_name = "input_ids"
_tied_weights_keys = {"lm_head.weight": "embeddings.word_embeddings.weight"}
def __init__(self, config):
super().__init__(config)
self.embeddings = BashkirRobertaEmbeddings(config)
self.layers = nn.ModuleList(
[BashkirRobertaLayer(config) for _ in range(config.num_hidden_layers)]
)
self.norm_final = nn.LayerNorm(config.hidden_size, eps=1e-12)
self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
self.post_init()
def get_input_embeddings(self):
return self.embeddings.word_embeddings
def set_input_embeddings(self, value):
self.embeddings.word_embeddings = value
def get_output_embeddings(self):
return self.lm_head
def set_output_embeddings(self, new_embeddings):
self.lm_head = new_embeddings
def tie_weights(self, **kwargs):
if self.config.tie_word_embeddings:
self.lm_head.weight = self.embeddings.word_embeddings.weight
def forward(self, input_ids, attention_mask=None, labels=None, return_dict=True, **kwargs):
if input_ids.shape[1] > self.config.max_position_embeddings:
raise ValueError(
f"Sequence length {input_ids.shape[1]} exceeds "
f"max_position_embeddings={self.config.max_position_embeddings}."
)
padding_mask = input_ids.eq(self.config.pad_token_id) if attention_mask is None else attention_mask.eq(0)
hidden_states = self.embeddings(input_ids)
for layer in self.layers:
hidden_states = layer(hidden_states, padding_mask=padding_mask)
hidden_states = self.norm_final(hidden_states)
logits = self.lm_head(hidden_states)
loss = None
if labels is not None:
loss = nn.functional.cross_entropy(
logits.view(-1, self.config.vocab_size), labels.view(-1), ignore_index=-100
)
if not return_dict:
return (loss, logits) if loss is not None else (logits,)
return MaskedLMOutput(loss=loss, logits=logits)
|