import math import torch import torch.nn as nn import torch.nn.functional as F from transformers import PreTrainedModel from transformers.modeling_outputs import ( MaskedLMOutput, BaseModelOutput, SequenceClassifierOutput, TokenClassifierOutput, QuestionAnsweringModelOutput, ) from .configuration_mybert import MyBertConfig def _build_rope_cache(head_dim, max_seq_len, base=10000.0): inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim)) t = torch.arange(max_seq_len, dtype=torch.float32) freqs = torch.outer(t, inv_freq) emb = torch.cat((freqs, freqs), dim=-1) return emb.cos(), emb.sin() def _rotate_half(x): x1, x2 = x.chunk(2, dim=-1) return torch.cat((-x2, x1), dim=-1) def _apply_rope(q, k, cos, sin): cos = cos.to(q.dtype)[None, None, :, :] sin = sin.to(q.dtype)[None, None, :, :] return (q * cos) + (_rotate_half(q) * sin), (k * cos) + (_rotate_half(k) * sin) class MyBertEmbeddings(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.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) self.dropout = nn.Dropout(config.hidden_dropout_prob) def forward(self, input_ids): return self.dropout(self.LayerNorm(self.word_embeddings(input_ids))) class MyBertSelfAttention(nn.Module): def __init__(self, config): super().__init__() self.num_attention_heads = config.num_attention_heads self.attention_head_size = config.hidden_size // config.num_attention_heads self.all_head_size = config.hidden_size b = config.use_bias self.query = nn.Linear(config.hidden_size, self.all_head_size, bias=b) self.key = nn.Linear(config.hidden_size, self.all_head_size, bias=b) self.value = nn.Linear(config.hidden_size, self.all_head_size, bias=b) self.dropout_prob = config.attention_probs_dropout_prob def forward(self, hidden_states, attention_mask=None, cos=None, sin=None): q, k, v = self.query(hidden_states), self.key(hidden_states), self.value(hidden_states) shp = q.size()[:-1] + (self.num_attention_heads, self.attention_head_size) q = q.view(*shp).transpose(1, 2) k = k.view(*shp).transpose(1, 2) v = v.view(*shp).transpose(1, 2) if cos is not None: q, k = _apply_rope(q, k, cos, sin) ctx = F.scaled_dot_product_attention( q, k, v, attn_mask=attention_mask, dropout_p=self.dropout_prob if self.training else 0.0, is_causal=False, ) ctx = ctx.transpose(1, 2).contiguous() return ctx.view(*ctx.size()[:-2], self.all_head_size) class MyBertSelfOutput(nn.Module): def __init__(self, config): super().__init__() self.dense = nn.Linear(config.hidden_size, config.hidden_size, bias=config.use_bias) self.dropout = nn.Dropout(config.hidden_dropout_prob) def forward(self, hidden_states): return self.dropout(self.dense(hidden_states)) class MyBertAttention(nn.Module): def __init__(self, config): super().__init__() self.self = MyBertSelfAttention(config) self.output = MyBertSelfOutput(config) def forward(self, hidden_states, attention_mask=None, cos=None, sin=None): return self.output(self.self(hidden_states, attention_mask, cos, sin)) class MyBertSwiGLU(nn.Module): def __init__(self, config): super().__init__() b = config.use_bias self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=b) self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=b) self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=b) self.dropout = nn.Dropout(config.hidden_dropout_prob) def forward(self, x): return self.dropout(self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))) class MyBertLayer(nn.Module): def __init__(self, config): super().__init__() self.attention_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) self.attention = MyBertAttention(config) self.ffn_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) self.mlp = MyBertSwiGLU(config) def forward(self, hidden_states, attention_mask=None, cos=None, sin=None): hidden_states = hidden_states + self.attention( self.attention_layernorm(hidden_states), attention_mask, cos, sin) hidden_states = hidden_states + self.mlp(self.ffn_layernorm(hidden_states)) return hidden_states class MyBertEncoder(nn.Module): def __init__(self, config): super().__init__() self.layer = nn.ModuleList([MyBertLayer(config) for _ in range(config.num_hidden_layers)]) self.final_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) def forward(self, hidden_states, attention_mask=None, cos=None, sin=None): for lyr in self.layer: hidden_states = lyr(hidden_states, attention_mask, cos, sin) return self.final_layernorm(hidden_states) class MyBertPredictionHeadTransform(nn.Module): def __init__(self, config): super().__init__() self.dense = nn.Linear(config.hidden_size, config.hidden_size) self.transform_act_fn = nn.GELU() self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) def forward(self, h): return self.LayerNorm(self.transform_act_fn(self.dense(h))) class MyBertLMPredictionHead(nn.Module): def __init__(self, config): super().__init__() self.transform = MyBertPredictionHeadTransform(config) self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=True) def forward(self, h): return self.decoder(self.transform(h)) class MyBertOnlyMLMHead(nn.Module): def __init__(self, config): super().__init__() self.predictions = MyBertLMPredictionHead(config) def forward(self, h): return self.predictions(h) class MyBertPreTrainedModel(PreTrainedModel): config_class = MyBertConfig base_model_prefix = "mybert" supports_gradient_checkpointing = False _no_split_modules = ["MyBertLayer"] def _init_weights(self, module): std = self.config.initializer_range if isinstance(module, nn.Linear): module.weight.data.normal_(mean=0.0, std=std) if module.bias is not None: module.bias.data.zero_() elif isinstance(module, nn.Embedding): module.weight.data.normal_(mean=0.0, std=std) if module.padding_idx is not None: module.weight.data[module.padding_idx].zero_() elif isinstance(module, nn.LayerNorm): module.bias.data.zero_() module.weight.data.fill_(1.0) if isinstance(module, MyBertModel): head_dim = self.config.hidden_size // self.config.num_attention_heads cos, sin = _build_rope_cache(head_dim, self.config.max_position_embeddings, self.config.rope_theta) if module.rope_cos.device.type != "meta": module.rope_cos.copy_(cos) module.rope_sin.copy_(sin) class MyBertModel(MyBertPreTrainedModel): def __init__(self, config): super().__init__(config) self.embeddings = MyBertEmbeddings(config) self.encoder = MyBertEncoder(config) head_dim = config.hidden_size // config.num_attention_heads cos, sin = _build_rope_cache(head_dim, config.max_position_embeddings, config.rope_theta) self.register_buffer("rope_cos", cos, persistent=False) self.register_buffer("rope_sin", sin, persistent=False) self.post_init() n = config.num_hidden_layers for lyr in self.encoder.layer: lyr.attention.output.dense.weight.data.normal_( 0.0, config.initializer_range / math.sqrt(2 * n)) lyr.mlp.down_proj.weight.data.normal_( 0.0, config.initializer_range / math.sqrt(2 * n)) def get_input_embeddings(self): return self.embeddings.word_embeddings def set_input_embeddings(self, v): self.embeddings.word_embeddings = v def forward(self, input_ids=None, attention_mask=None, return_dict=True, **kw): T = input_ids.shape[1] if T > self.rope_cos.shape[0]: raise ValueError(f"seq_len {T} > max_position_embeddings {self.rope_cos.shape[0]}") cos, sin = self.rope_cos[:T], self.rope_sin[:T] attn_mask = attention_mask.bool()[:, None, None, :] if attention_mask is not None else None seq = self.encoder(self.embeddings(input_ids), attn_mask, cos, sin) if not return_dict: return (seq,) return BaseModelOutput(last_hidden_state=seq) class MyBertForMaskedLM(MyBertPreTrainedModel): _tied_weights_keys = { "cls.predictions.decoder.weight": "mybert.embeddings.word_embeddings.weight", } def __init__(self, config): super().__init__(config) self.mybert = MyBertModel(config) self.cls = MyBertOnlyMLMHead(config) self.post_init() def get_output_embeddings(self): return self.cls.predictions.decoder def set_output_embeddings(self, new): self.cls.predictions.decoder = new def forward(self, input_ids=None, attention_mask=None, labels=None, label_smoothing=0.0, return_dict=True, **kw): seq = self.mybert(input_ids=input_ids, attention_mask=attention_mask, return_dict=True).last_hidden_state if labels is not None and self.config.sparse_prediction: flat_h = seq.view(-1, seq.size(-1)) flat_y = labels.view(-1) idx = (flat_y != -100).nonzero(as_tuple=True)[0] logits = self.cls(flat_h.index_select(0, idx)) loss = F.cross_entropy(logits.float(), flat_y.index_select(0, idx), label_smoothing=label_smoothing) return MaskedLMOutput(loss=loss, logits=None) logits = self.cls(seq) loss = None if labels is not None: loss = F.cross_entropy(logits.float().view(-1, self.config.vocab_size), labels.view(-1), ignore_index=-100, label_smoothing=label_smoothing) if not return_dict: return ((loss, logits) if loss is not None else (logits,)) return MaskedLMOutput(loss=loss, logits=logits) class MyBertForSequenceClassification(MyBertPreTrainedModel): def __init__(self, config): super().__init__(config) self.num_labels = config.num_labels self.config = config self.mybert = MyBertModel(config) classifier_dropout = ( config.classifier_dropout if getattr(config, "classifier_dropout", None) is not None else config.hidden_dropout_prob ) self.dropout = nn.Dropout(classifier_dropout) self.classifier = nn.Linear(config.hidden_size, config.num_labels) self.post_init() def forward(self, input_ids=None, attention_mask=None, labels=None, return_dict=True, **kw): seq = self.mybert(input_ids=input_ids, attention_mask=attention_mask, return_dict=True).last_hidden_state pooled = seq[:, 0] logits = self.classifier(self.dropout(pooled)) loss = None if labels is not None: if self.config.problem_type is None: if self.num_labels == 1: self.config.problem_type = "regression" elif self.num_labels > 1 and labels.dtype in (torch.long, torch.int): self.config.problem_type = "single_label_classification" else: self.config.problem_type = "multi_label_classification" if self.config.problem_type == "regression": loss_fct = nn.MSELoss() loss = (loss_fct(logits.squeeze(), labels.squeeze()) if self.num_labels == 1 else loss_fct(logits, labels)) elif self.config.problem_type == "single_label_classification": loss = nn.CrossEntropyLoss()(logits.view(-1, self.num_labels), labels.view(-1)) else: loss = nn.BCEWithLogitsLoss()(logits, labels) if not return_dict: return ((loss, logits) if loss is not None else (logits,)) return SequenceClassifierOutput(loss=loss, logits=logits) class MyBertForTokenClassification(MyBertPreTrainedModel): def __init__(self, config): super().__init__(config) self.num_labels = config.num_labels self.mybert = MyBertModel(config) classifier_dropout = ( config.classifier_dropout if getattr(config, "classifier_dropout", None) is not None else config.hidden_dropout_prob ) self.dropout = nn.Dropout(classifier_dropout) self.classifier = nn.Linear(config.hidden_size, config.num_labels) self.post_init() def forward(self, input_ids=None, attention_mask=None, labels=None, return_dict=True, **kw): seq = self.mybert(input_ids=input_ids, attention_mask=attention_mask, return_dict=True).last_hidden_state logits = self.classifier(self.dropout(seq)) loss = None if labels is not None: loss = nn.CrossEntropyLoss()(logits.view(-1, self.num_labels), labels.view(-1)) if not return_dict: return ((loss, logits) if loss is not None else (logits,)) return TokenClassifierOutput(loss=loss, logits=logits) class MyBertForQuestionAnswering(MyBertPreTrainedModel): def __init__(self, config): super().__init__(config) self.num_labels = 2 self.mybert = MyBertModel(config) self.qa_outputs = nn.Linear(config.hidden_size, 2) self.post_init() def forward(self, input_ids=None, attention_mask=None, start_positions=None, end_positions=None, return_dict=True, **kw): seq = self.mybert(input_ids=input_ids, attention_mask=attention_mask, return_dict=True).last_hidden_state start_logits, end_logits = self.qa_outputs(seq).split(1, dim=-1) start_logits = start_logits.squeeze(-1).contiguous() end_logits = end_logits.squeeze(-1).contiguous() total_loss = None if start_positions is not None and end_positions is not None: if start_positions.dim() > 1: start_positions = start_positions.squeeze(-1) if end_positions.dim() > 1: end_positions = end_positions.squeeze(-1) ignored_index = start_logits.size(1) start_positions = start_positions.clamp(0, ignored_index) end_positions = end_positions.clamp(0, ignored_index) loss_fct = nn.CrossEntropyLoss(ignore_index=ignored_index) total_loss = (loss_fct(start_logits, start_positions) + loss_fct(end_logits, end_positions)) / 2 if not return_dict: out = (start_logits, end_logits) return ((total_loss,) + out) if total_loss is not None else out return QuestionAnsweringModelOutput( loss=total_loss, start_logits=start_logits, end_logits=end_logits)