modeling_mybert.py

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)