"""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, )