""" RoFormer Model Architecture. This module defines a custom Transformer model based on the non-autoregressive architecture, adapted for ancient Slavic texts. It overrides the standard HuggingFace BERT encoder to introduce the following innovations: 1. Pure Character-Level Representation: Better handle highly inflected languages and physically fragmented text (e.g., birch bark). 2. Rotary Position Embeddings (RoPE): Replaces absolute positional embeddings. RoPE rotates query and key vectors in the self-attention mechanism, allowing the model to natively understand relative distances between characters across variable-length gaps. 3. Multi-Task Learning (4 Heads): - Restore Head (Local): Standard Masked Language Modeling (MLM) to predict missing characters `[-]`. - Unk Head (Local): A binary classifier predicting whether a gap token `[#]` should be expanded. - Date Head (Global): Predicts the document's chronological distribution (reads from `[SOS]`). - Region Head (Global): Classifies the dialect/region of the document (reads from `[SOS]`). """ import math from dataclasses import dataclass from typing import Optional, Tuple import torch import torch.nn as nn try: from .config import RoFormerConfig # package-style import (e.g. inference.py) except ImportError: from config import RoFormerConfig # script-mode import (e.g. train.py) from transformers import BertPreTrainedModel from transformers.models.bert.modeling_bert import BertEncoder, BertSelfAttention from transformers.utils.generic import ModelOutput from transformers.models.llama.modeling_llama import LlamaRotaryEmbedding, apply_rotary_pos_emb from transformers.models.llama.configuration_llama import LlamaConfig # CUSTOM ATTENTION WITH RoPE class BertSelfAttentionWithRoPE(BertSelfAttention): def __init__(self, config: RoFormerConfig): super().__init__(config) # Some transformers versions' BertSelfAttention.__init__ doesn't keep # a `self.config` reference -- set it explicitly so forward() below # doesn't depend on that varying by version. self.config = config llama_config = LlamaConfig( hidden_size=self.attention_head_size * self.num_attention_heads, num_attention_heads=self.num_attention_heads, max_position_embeddings=config.max_position_embeddings, ) self.rotary_emb = LlamaRotaryEmbedding( config=llama_config, ) def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor: new_x_shape = x.size()[:-1] + ( self.num_attention_heads, self.attention_head_size, ) x = x.view(new_x_shape) return x.permute(0, 2, 1, 3) def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, head_mask: Optional[torch.Tensor] = None, encoder_hidden_states: Optional[torch.Tensor] = None, encoder_attention_mask: Optional[torch.Tensor] = None, past_key_value: Optional[Tuple[torch.Tensor]] = None, output_attentions: bool = False, **kwargs, ) -> Tuple[torch.Tensor, ...]: """ Forward pass for the custom Self-Attention layer injecting RoPE. Args: hidden_states (torch.Tensor): Input hidden states. attention_mask (Optional[torch.Tensor]): Mask to avoid performing attention on padding token indices. head_mask (Optional[torch.Tensor]): Mask to nullify selected heads of the self-attention modules. encoder_hidden_states (Optional[torch.Tensor]): Used in cross-attention (not applicable here). encoder_attention_mask (Optional[torch.Tensor]): Used in cross-attention (not applicable here). past_key_value (Optional[Tuple[torch.Tensor]]): Cached key and value states. output_attentions (bool): Whether or not to return the attentions tensors. """ output_attentions = getattr( self.config, "output_attentions", output_attentions ) or kwargs.get("output_attentions", output_attentions) mixed_query_layer = self.query(hidden_states) key_layer = self.transpose_for_scores(self.key(hidden_states)) value_layer = self.transpose_for_scores(self.value(hidden_states)) query_layer = self.transpose_for_scores(mixed_query_layer) # Apply RoPE (Rotary Position Embeddings) to Query and Key matrices seq_length = hidden_states.shape[1] position_ids = torch.arange( seq_length, dtype=torch.long, device=hidden_states.device ) position_ids = position_ids.unsqueeze(0).expand(hidden_states.shape[0], -1) cos, sin = self.rotary_emb(value_layer, position_ids) query_layer, key_layer = apply_rotary_pos_emb( query_layer, key_layer, cos, sin ) # Manual attention instead of SDPA: ROCm's SDPA backward pass produces # nan/inf gradients in bfloat16 once LR ramps up past warmup (~600 # steps) -- confirmed twice on MI300X (see commits a283f5f/9e8a617). # Do not switch this back to SDPA even for the non-output_attentions # path; the regression that reintroduced it (aa11bfa) was a mistake. attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2)) attention_scores = attention_scores / math.sqrt(self.attention_head_size) if attention_mask is not None: attention_scores = attention_scores + attention_mask attention_probs = nn.functional.softmax(attention_scores, dim=-1) attention_probs = self.dropout(attention_probs) if output_attentions: self._last_attention = attention_probs else: self._last_attention = None context_layer = torch.matmul(attention_probs, value_layer) context_layer = context_layer.permute(0, 2, 1, 3).contiguous() new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,) context_layer = context_layer.view(*new_context_layer_shape) return context_layer, (attention_probs if output_attentions else None) @dataclass class RoFormerOutput(ModelOutput): logits_restore: torch.FloatTensor = None logits_unk: torch.FloatTensor = None logits_date: Optional[torch.FloatTensor] = None logits_region: Optional[torch.FloatTensor] = None attentions: Optional[Tuple[torch.FloatTensor]] = None # MAIN MODEL ARCHITECTURE class RoFormer(BertPreTrainedModel): config_class = RoFormerConfig def __init__( self, config: RoFormerConfig, num_date_bins: int = 20, num_regions: int = 4 ): super().__init__(config) # Disable absolute positional embeddings inherited from BERT config.position_embedding_type = "none" # 1. Pure Character Embeddings Layer (No Word IDs) self.char_embeddings = nn.Embedding( config.vocab_char_size, config.hidden_size, padding_idx=config.pad_token_id ) self.emb_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) self.emb_dropout = nn.Dropout(config.hidden_dropout_prob) # 2. Transformer Encoder with RoPE Injection self.encoder = BertEncoder(config) for layer in self.encoder.layer: # Overwrite the standard self-attention module with our custom RoPE implementation layer.attention.self = BertSelfAttentionWithRoPE(config) # 3. Multi-Task Output Heads # Local Head A: Restores missing characters (tied to char_embeddings weights) self.restore_dense = nn.Linear(config.hidden_size, config.hidden_size) self.restore_act = nn.GELU() self.restore_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) self.restore_bias = nn.Parameter(torch.zeros(config.vocab_char_size)) # Local Head B: Predicts gap extension for the `[#]` token self.unk_head = nn.Sequential( nn.Linear(config.hidden_size, config.hidden_size), nn.GELU(), nn.Linear(config.hidden_size, 2), ) # Global Head A: Predicts historical date distribution (Reads from [SOS]) self.date_head = nn.Sequential( nn.Linear(config.hidden_size, config.hidden_size), nn.GELU(), nn.Linear(config.hidden_size, num_date_bins), ) # Global Head B: Classifies dialect/region (Reads from [SOS]) self.region_head = nn.Sequential( nn.Linear(config.hidden_size, config.hidden_size), nn.GELU(), nn.Linear(config.hidden_size, num_regions), ) self.post_init() def get_input_embeddings(self) -> nn.Embedding: """ Required by HuggingFace to support `resize_token_embeddings`. Returns: nn.Embedding: The embedding layer. """ return self.char_embeddings def set_input_embeddings(self, value: nn.Embedding) -> None: """ Required by HuggingFace to support `resize_token_embeddings`. Args: value (nn.Embedding): The new embedding layer. """ self.char_embeddings = value def forward( self, input_ids: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, output_attentions: Optional[bool] = None, **kwargs, ) -> RoFormerOutput: """ Forward pass for the multi-task Transformer. Args: input_ids (Optional[torch.Tensor]): Indices of input sequence tokens in the vocabulary. attention_mask (Optional[torch.Tensor]): Mask to avoid performing attention on padding token indices. output_attentions (Optional[bool]): Whether or not to return the attentions tensors. **kwargs: Additional keyword arguments. Returns: RoFormerOutput: A custom dataclass containing 4 sets of logits and optional attention weights. """ output_attentions = ( output_attentions if output_attentions is not None else self.config.output_attentions ) # Fast character embedding generation x = self.char_embeddings(input_ids) x = self.emb_norm(x) x = self.emb_dropout(x) ext_mask = self.get_extended_attention_mask(attention_mask, input_ids.shape) # Pass through the RoPE-enabled encoder torso enc_out = self.encoder( x, attention_mask=ext_mask, output_attentions=output_attentions, return_dict=True, ) seq = enc_out.last_hidden_state # Shape: [Batch, Seq_Len, Hidden_Dim] # Process local sequence heads (applied to all tokens) x_res = self.restore_norm(self.restore_act(self.restore_dense(seq))) # Tied-embedding projection. Aeneas pairs sqrt(hidden_size) normalization + # dropout-on-the-weight-matrix with weight_decay=0.0 and LAMB (see config_greek.py); # under plain AdamW + weight_decay>0 on a short from-scratch budget, either half of # that combination may itself cap restore-head confidence (ablated independently of # weight_decay -- see RoFormerConfig.tied_proj_scale/tied_proj_dropout). emb_weights = self.char_embeddings.weight if self.config.tied_proj_dropout: emb_weights = self.emb_dropout(emb_weights) logits_restore = x_res @ emb_weights.T if self.config.tied_proj_scale: logits_restore = logits_restore / math.sqrt(emb_weights.shape[-1]) logits_restore = logits_restore + self.restore_bias logits_unk = self.unk_head(seq) # Process global heads (applied exclusively to the [SOS] token at index 0) sos_vectors = seq[:, 0, :] logits_date = self.date_head(sos_vectors) logits_region = self.region_head(sos_vectors) # Manually extract the attention from the last layer to bypass HF tuple dropping last_attn = self.encoder.layer[-1].attention.self._last_attention extracted_attentions = (last_attn,) if last_attn is not None else None return RoFormerOutput( logits_restore=logits_restore, logits_unk=logits_unk, logits_date=logits_date, logits_region=logits_region, attentions=extracted_attentions, )