modeling_mathbananamind.py

"""MathBananaMind model implementation for Hugging Face Transformers."""

from typing import Optional

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel
from transformers.generation.utils import GenerationMixin
from transformers.modeling_outputs import CausalLMOutputWithPast

from .configuration_mathbananamind import MathBananaMindConfig


class RMSNorm(nn.Module):
    def __init__(self, width, eps):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(width))
        self.eps = eps

    def forward(self, hidden):
        source_dtype = hidden.dtype
        hidden_float = hidden.float()
        normalized = hidden_float * torch.rsqrt(
            hidden_float.pow(2).mean(dim=-1, keepdim=True) + self.eps
        )
        return (normalized * self.weight.float()).to(source_dtype)


def rotate_half(hidden):
    first, second = hidden.chunk(2, dim=-1)
    return torch.cat((-second, first), dim=-1)


class RotaryEmbedding(nn.Module):
    def __init__(self, head_dim, max_seq_len, theta):
        super().__init__()
        self.head_dim = head_dim
        self.max_seq_len = max_seq_len
        self.theta = theta
        self._cosine_cache = None
        self._sine_cache = None

    def _get_cache(self, device):
        if (
            self._cosine_cache is None
            or self._cosine_cache.device != device
        ):
            # Build this outside __init__ so from_pretrained's meta-device
            # initialization cannot leave non-persistent buffers unmaterialized.
            head_dim = self.head_dim
            theta = self.theta
            max_seq_len = self.max_seq_len
            inverse_frequency = 1.0 / (
                theta
                ** (
                    torch.arange(0, head_dim, 2, dtype=torch.float32)
                    / head_dim
                )
            )
            positions = torch.arange(max_seq_len, dtype=torch.float32)
            frequencies = torch.outer(positions, inverse_frequency)
            embedding = torch.cat((frequencies, frequencies), dim=-1)
            self._cosine_cache = embedding.cos().to(device)
            self._sine_cache = embedding.sin().to(device)
        return self._cosine_cache, self._sine_cache

    def forward(self, query, key):
        sequence_length = query.shape[-2]
        cosine_cache, sine_cache = self._get_cache(query.device)
        cosine = cosine_cache[:sequence_length].to(query.dtype)[None, None, :, :]
        sine = sine_cache[:sequence_length].to(query.dtype)[None, None, :, :]
        return (
            query * cosine + rotate_half(query) * sine,
            key * cosine + rotate_half(key) * sine,
        )


class CausalSelfAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.num_heads = config.num_attention_heads
        self.num_kv_heads = config.num_key_value_heads
        self.head_dim = config.head_dim
        self.kv_width = self.num_kv_heads * self.head_dim
        self.dropout = config.attention_dropout
        self.query = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
        self.key = nn.Linear(config.hidden_size, self.kv_width, bias=False)
        self.value = nn.Linear(config.hidden_size, self.kv_width, bias=False)
        self.output = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
        self.rope = RotaryEmbedding(
            self.head_dim,
            config.max_position_embeddings,
            config.rope_theta,
        )

    def forward(self, hidden, attention_mask=None):
        batch_size, sequence_length, _ = hidden.shape
        query = self.query(hidden).view(
            batch_size, sequence_length, self.num_heads, self.head_dim
        ).transpose(1, 2)
        key = self.key(hidden).view(
            batch_size, sequence_length, self.num_kv_heads, self.head_dim
        ).transpose(1, 2)
        value = self.value(hidden).view(
            batch_size, sequence_length, self.num_kv_heads, self.head_dim
        ).transpose(1, 2)
        query, key = self.rope(query, key)
        repeats = self.num_heads // self.num_kv_heads
        if repeats > 1:
            key = key.repeat_interleave(repeats, dim=1)
            value = value.repeat_interleave(repeats, dim=1)

        sdpa_mask = None
        is_causal = True
        if attention_mask is not None:
            key_mask = attention_mask.to(torch.bool)[:, None, None, :]
            causal_mask = torch.ones(
                sequence_length,
                sequence_length,
                dtype=torch.bool,
                device=hidden.device,
            ).tril()
            sdpa_mask = key_mask & causal_mask[None, None, :, :]
            is_causal = False

        attended = F.scaled_dot_product_attention(
            query,
            key,
            value,
            attn_mask=sdpa_mask,
            dropout_p=self.dropout if self.training else 0.0,
            is_causal=is_causal,
        )
        attended = attended.transpose(1, 2).contiguous().view(
            batch_size, sequence_length, self.num_heads * self.head_dim
        )
        return self.output(attended)


class SwiGLU(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.gate = nn.Linear(
            config.hidden_size, config.intermediate_size, bias=False
        )
        self.up = nn.Linear(
            config.hidden_size, config.intermediate_size, bias=False
        )
        self.down = nn.Linear(
            config.intermediate_size, config.hidden_size, bias=False
        )

    def forward(self, hidden):
        return self.down(F.silu(self.gate(hidden)) * self.up(hidden))


class DecoderBlock(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.attention_norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
        self.attention = CausalSelfAttention(config)
        self.mlp_norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
        self.mlp = SwiGLU(config)

    def forward(self, hidden, attention_mask=None):
        hidden = hidden + self.attention(
            self.attention_norm(hidden), attention_mask=attention_mask
        )
        hidden = hidden + self.mlp(self.mlp_norm(hidden))
        return hidden


class MathBananaMindPreTrainedModel(PreTrainedModel):
    config_class = MathBananaMindConfig
    base_model_prefix = ""
    supports_gradient_checkpointing = False
    _supports_cache_class = False

    def _init_weights(self, module):
        if isinstance(module, (nn.Linear, nn.Embedding)):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)


class MathBananaMindForCausalLM(MathBananaMindPreTrainedModel, GenerationMixin):
    def __init__(self, config):
        super().__init__(config)
        self.token_embedding = nn.Embedding(
            config.vocab_size, config.embedding_size
        )
        self.input_projection = nn.Linear(
            config.embedding_size, config.hidden_size, bias=False
        )
        self.output_projection = nn.Linear(
            config.hidden_size, config.embedding_size, bias=False
        )
        self.blocks = nn.ModuleList(
            DecoderBlock(config) for _ in range(config.num_hidden_layers)
        )
        self.final_norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
        self.post_init()

    def get_input_embeddings(self):
        return self.token_embedding

    def set_input_embeddings(self, value):
        self.token_embedding = value

    def get_output_embeddings(self):
        return None

    def prepare_inputs_for_generation(
        self,
        input_ids,
        past_key_values=None,
        attention_mask=None,
        **kwargs,
    ):
        return {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "use_cache": False,
        }

    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        labels=None,
        past_key_values: Optional[object] = None,
        use_cache=False,
        **kwargs,
    ):
        if input_ids is None:
            raise ValueError("input_ids must be provided")
        if input_ids.shape[1] > self.config.max_position_embeddings:
            raise ValueError("input sequence exceeds max_position_embeddings")

        hidden = self.input_projection(self.token_embedding(input_ids))
        for block in self.blocks:
            hidden = block(hidden, attention_mask=attention_mask)
        hidden = self.final_norm(hidden)
        token_hidden = self.output_projection(hidden)
        logits = F.linear(token_hidden, self.token_embedding.weight)

        loss = None
        if labels is not None:
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            loss = F.cross_entropy(
                shift_logits.float().reshape(-1, self.config.vocab_size),
                shift_labels.reshape(-1),
                ignore_index=-100,
            )

        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=None,
        )