Update modeling_hf_nomic_bert.py

modeling_hf_nomic_bert.py

@@ -16,7 +16,7 @@ from einops import rearrange, repeat
 from transformers import GPT2Config, PreTrainedModel
 from transformers.models.bert.modeling_bert import (
     BaseModelOutputWithPoolingAndCrossAttentions,
-    BertForPreTrainingOutput,
+    MaskedLMOutput,
     SequenceClassifierOutput
 )
 
@@ -323,6 +323,8 @@ class NomicBertPreTrainedModel(PreTrainedModel):
         rotary_scaling_factor = kwargs.pop("rotary_scaling_factor", None)
         if rotary_scaling_factor:
             config.rotary_scaling_factor = rotary_scaling_factor
+        else:
+            config.rotary_scaling_factor = None
         if config.n_positions <= 0 and config.rotary_emb_fraction > 0:
             config.n_positions = 2048
         if num_labels:
@@ -1145,9 +1147,11 @@ class NomicBertForPreTraining(NomicBertPreTrainedModel):
             )
             total_loss = masked_lm_loss.float()
 
-        return BertForPreTrainingOutput(
+        return MaskedLMOutput(
             loss=total_loss,
-            prediction_logits=prediction_scores,
+            logits=prediction_scores,
+            hidden_states=outputs.hidden_states,
+            attentions=None, 
         )