Upload modeling_baseline.py with huggingface_hub

modeling_baseline.py

@@ -2,16 +2,11 @@
 import torch
 import torch.nn as nn
 from transformers import PreTrainedModel
-
-try:
-    from .configuration_baseline import BaselineConfig
-except ImportError:
-    from configuration_baseline import BaselineConfig
-
+from .configuration_baseline import BaselineConfig
 try:
     from x_transformers import TransformerWrapper, Encoder
 except ImportError:
-    raise ImportError("To use this model, you must run: pip install x-transformers")
+    raise ImportError("pip install x-transformers")
 
 class BaselineModel(PreTrainedModel):
     config_class = BaselineConfig
@@ -19,13 +14,11 @@ class BaselineModel(PreTrainedModel):
     def __init__(self, config):
         super().__init__(config)
         self.config = config
-        
-        # Instantiate the x_transformers model
         self.model = TransformerWrapper(
             num_tokens=config.vocab_size,
             max_seq_len=config.seq_len,
-            use_abs_pos_emb=False,    # RoPE is enabled in Encoder
-            tie_embedding=True,       # Matches training
+            use_abs_pos_emb=False,
+            tie_embedding=True,
             attn_layers=Encoder(
                 dim=config.d_model,
                 depth=config.depth,
@@ -33,21 +26,19 @@ class BaselineModel(PreTrainedModel):
                 layer_dropout=config.dropout,
                 attn_dropout=config.dropout,
                 ff_dropout=config.dropout,
-                rotary_pos_emb=True,  # Matches training
-                attn_flash=True,      # Matches training
+                rotary_pos_emb=True,
+                attn_flash=True,
                 use_scalenorm=False
             )
         )
+        # TIE FIX
+        if hasattr(self.model.token_emb, 'emb'):
+            self.model.to_logits.weight = self.model.token_emb.emb.weight
+        else:
+            self.model.to_logits.weight = self.model.token_emb.weight
 
     def forward(self, input_ids, labels=None, mask=None):
-        # x_transformers takes 'mask' argument if provided
         logits = self.model(input_ids, mask=mask)
-        
-        loss = None
         if labels is not None:
-            loss_fct = nn.CrossEntropyLoss()
-            # Reshape for loss calculation
-            loss = loss_fct(logits.view(-1, self.config.vocab_size), labels.view(-1))
-            return {"loss": loss, "logits": logits}
-
+            return {"loss": nn.CrossEntropyLoss()(logits.view(-1, self.config.vocab_size), labels.view(-1)), "logits": logits}
         return logits