Fix model code

modeling_blaser.py

@@ -104,12 +104,14 @@ class BlaserCore(nn.Module):
 
 
 # ---------------- HF MODEL WRAPPER ---------------- #
+
 class BlaserModel(PreTrainedModel):
     config_class = BlaserConfig
 
     def __init__(self, config: BlaserConfig):
         super().__init__(config)
-        self.core = BlaserCore(
+        # Instead of self.core, assign directly
+        self.mlp = BlaserCore(
             embedding_dim=config.embedding_dim,
             output_dim=config.output_dim,
             hidden_dims=config.hidden_dims,
@@ -118,7 +120,14 @@ class BlaserModel(PreTrainedModel):
             input_form=config.input_form,
             norm_emb=config.norm_emb,
             output_act=config.output_act,
-        )
+        ).mlp  # only take the Sequential MLP
 
     def forward(self, src, mt, ref=None):
-        return self.core(src, mt, ref)
+        # The old checkpoint expects the input feature processing inside BlaserCore
+        proc = BlaserCore._featurize(
+            self.mlp,  # pass self as `self` for static call
+            src=BlaserCore._norm(self.mlp, src),
+            mt=BlaserCore._norm(self.mlp, mt),
+            ref=BlaserCore._norm(self.mlp, ref)
+        )
+        return self.mlp(proc)