Fix dynamic import of configuration_bashkir_roberta

modeling_bashkir_roberta.py

@@ -1,114 +1,117 @@
-"""Exact Hugging Face-compatible implementation of BashkirRoBERTa Pre-LN."""
-
-import torch
-from torch import nn
-from transformers import PreTrainedModel
-from transformers.modeling_outputs import MaskedLMOutput
-
-from configuration_bashkir_roberta import BashkirRobertaConfig
-
-
-class BashkirRobertaEmbeddings(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.position_embeddings = nn.Embedding(
-            config.max_position_embeddings, config.hidden_size
-        )
-        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=1e-12)
-        self.dropout = nn.Dropout(config.hidden_dropout_prob)
-
-    def forward(self, input_ids):
-        sequence_length = input_ids.shape[1]
-        positions = torch.arange(
-            sequence_length, dtype=torch.long, device=input_ids.device
-        ).unsqueeze(0)
-        hidden_states = self.word_embeddings(input_ids) + self.position_embeddings(positions)
-        return self.dropout(self.LayerNorm(hidden_states))
-
-
-class BashkirRobertaLayer(nn.Module):
-    """Pre-LayerNorm transformer block, matching ``train_flagship.py`` exactly."""
-
-    def __init__(self, config):
-        super().__init__()
-        self.self_attn = nn.MultiheadAttention(
-            config.hidden_size,
-            config.num_attention_heads,
-            dropout=config.hidden_dropout_prob,
-            batch_first=True,
-        )
-        self.norm1 = nn.LayerNorm(config.hidden_size, eps=1e-12)
-        self.ffn = nn.Sequential(
-            nn.Linear(config.hidden_size, config.intermediate_size),
-            nn.GELU(),
-            nn.Dropout(config.hidden_dropout_prob),
-            nn.Linear(config.intermediate_size, config.hidden_size),
-            nn.Dropout(config.hidden_dropout_prob),
-        )
-        self.norm2 = nn.LayerNorm(config.hidden_size, eps=1e-12)
-
-    def forward(self, hidden_states, padding_mask=None):
-        normalized = self.norm1(hidden_states)
-        attention, _ = self.self_attn(
-            normalized, normalized, normalized, key_padding_mask=padding_mask
-        )
-        hidden_states = hidden_states + attention
-        return hidden_states + self.ffn(self.norm2(hidden_states))
-
-
-class BashkirRobertaForMaskedLM(PreTrainedModel):
-    config_class = BashkirRobertaConfig
-    base_model_prefix = "bashkir_roberta"
-    main_input_name = "input_ids"
-    _tied_weights_keys = {"lm_head.weight": "embeddings.word_embeddings.weight"}
-
-    def __init__(self, config):
-        super().__init__(config)
-        self.embeddings = BashkirRobertaEmbeddings(config)
-        self.layers = nn.ModuleList(
-            [BashkirRobertaLayer(config) for _ in range(config.num_hidden_layers)]
-        )
-        self.norm_final = nn.LayerNorm(config.hidden_size, eps=1e-12)
-        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
-        self.post_init()
-
-    def get_input_embeddings(self):
-        return self.embeddings.word_embeddings
-
-    def set_input_embeddings(self, value):
-        self.embeddings.word_embeddings = value
-
-    def get_output_embeddings(self):
-        return self.lm_head
-
-    def set_output_embeddings(self, new_embeddings):
-        self.lm_head = new_embeddings
-
-    def tie_weights(self, **kwargs):
-        if self.config.tie_word_embeddings:
-            self.lm_head.weight = self.embeddings.word_embeddings.weight
-
-    def forward(self, input_ids, attention_mask=None, labels=None, return_dict=True, **kwargs):
-        if input_ids.shape[1] > self.config.max_position_embeddings:
-            raise ValueError(
-                f"Sequence length {input_ids.shape[1]} exceeds "
-                f"max_position_embeddings={self.config.max_position_embeddings}."
-            )
-        padding_mask = input_ids.eq(self.config.pad_token_id) if attention_mask is None else attention_mask.eq(0)
-        hidden_states = self.embeddings(input_ids)
-        for layer in self.layers:
-            hidden_states = layer(hidden_states, padding_mask=padding_mask)
-        hidden_states = self.norm_final(hidden_states)
-        logits = self.lm_head(hidden_states)
-
-        loss = None
-        if labels is not None:
-            loss = nn.functional.cross_entropy(
-                logits.view(-1, self.config.vocab_size), labels.view(-1), ignore_index=-100
-            )
-        if not return_dict:
-            return (loss, logits) if loss is not None else (logits,)
-        return MaskedLMOutput(loss=loss, logits=logits)
+"""Exact Hugging Face-compatible implementation of BashkirRoBERTa Pre-LN."""
+
+import torch
+from torch import nn
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import MaskedLMOutput
+
+try:
+    from .configuration_bashkir_roberta import BashkirRobertaConfig
+except (ImportError, ValueError):
+    from configuration_bashkir_roberta import BashkirRobertaConfig
+
+
+class BashkirRobertaEmbeddings(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.position_embeddings = nn.Embedding(
+            config.max_position_embeddings, config.hidden_size
+        )
+        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=1e-12)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+    def forward(self, input_ids):
+        sequence_length = input_ids.shape[1]
+        positions = torch.arange(
+            sequence_length, dtype=torch.long, device=input_ids.device
+        ).unsqueeze(0)
+        hidden_states = self.word_embeddings(input_ids) + self.position_embeddings(positions)
+        return self.dropout(self.LayerNorm(hidden_states))
+
+
+class BashkirRobertaLayer(nn.Module):
+    """Pre-LayerNorm transformer block, matching ``train_flagship.py`` exactly."""
+
+    def __init__(self, config):
+        super().__init__()
+        self.self_attn = nn.MultiheadAttention(
+            config.hidden_size,
+            config.num_attention_heads,
+            dropout=config.hidden_dropout_prob,
+            batch_first=True,
+        )
+        self.norm1 = nn.LayerNorm(config.hidden_size, eps=1e-12)
+        self.ffn = nn.Sequential(
+            nn.Linear(config.hidden_size, config.intermediate_size),
+            nn.GELU(),
+            nn.Dropout(config.hidden_dropout_prob),
+            nn.Linear(config.intermediate_size, config.hidden_size),
+            nn.Dropout(config.hidden_dropout_prob),
+        )
+        self.norm2 = nn.LayerNorm(config.hidden_size, eps=1e-12)
+
+    def forward(self, hidden_states, padding_mask=None):
+        normalized = self.norm1(hidden_states)
+        attention, _ = self.self_attn(
+            normalized, normalized, normalized, key_padding_mask=padding_mask
+        )
+        hidden_states = hidden_states + attention
+        return hidden_states + self.ffn(self.norm2(hidden_states))
+
+
+class BashkirRobertaForMaskedLM(PreTrainedModel):
+    config_class = BashkirRobertaConfig
+    base_model_prefix = "bashkir_roberta"
+    main_input_name = "input_ids"
+    _tied_weights_keys = {"lm_head.weight": "embeddings.word_embeddings.weight"}
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.embeddings = BashkirRobertaEmbeddings(config)
+        self.layers = nn.ModuleList(
+            [BashkirRobertaLayer(config) for _ in range(config.num_hidden_layers)]
+        )
+        self.norm_final = nn.LayerNorm(config.hidden_size, eps=1e-12)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embeddings.word_embeddings
+
+    def set_input_embeddings(self, value):
+        self.embeddings.word_embeddings = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def tie_weights(self, **kwargs):
+        if self.config.tie_word_embeddings:
+            self.lm_head.weight = self.embeddings.word_embeddings.weight
+
+    def forward(self, input_ids, attention_mask=None, labels=None, return_dict=True, **kwargs):
+        if input_ids.shape[1] > self.config.max_position_embeddings:
+            raise ValueError(
+                f"Sequence length {input_ids.shape[1]} exceeds "
+                f"max_position_embeddings={self.config.max_position_embeddings}."
+            )
+        padding_mask = input_ids.eq(self.config.pad_token_id) if attention_mask is None else attention_mask.eq(0)
+        hidden_states = self.embeddings(input_ids)
+        for layer in self.layers:
+            hidden_states = layer(hidden_states, padding_mask=padding_mask)
+        hidden_states = self.norm_final(hidden_states)
+        logits = self.lm_head(hidden_states)
+
+        loss = None
+        if labels is not None:
+            loss = nn.functional.cross_entropy(
+                logits.view(-1, self.config.vocab_size), labels.view(-1), ignore_index=-100
+            )
+        if not return_dict:
+            return (loss, logits) if loss is not None else (logits,)
+        return MaskedLMOutput(loss=loss, logits=logits)