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
 )
 
@@ -321,7 +321,10 @@ class NomicBertPreTrainedModel(PreTrainedModel):
         ignore_mismatched_shapes = kwargs.pop("ignore_mismatched_sizes", False)
         num_labels = kwargs.pop("num_labels", None)
         rotary_scaling_factor = kwargs.pop("rotary_scaling_factor", None)
-        config.rotary_scaling_factor = rotary_scaling_factor
+        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:
@@ -330,7 +333,10 @@ class NomicBertPreTrainedModel(PreTrainedModel):
         if "add_pooling_layer" in kwargs:
             model = cls(config, *inputs, add_pooling_layer=kwargs.pop("add_pooling_layer"))
         else:
-            model = cls(config, *inputs, add_pooling_layer=False)
+            if cls == NomicBertModel:
+                model = cls(config, *inputs, add_pooling_layer=False)
+            else:
+                model = cls(config, *inputs)
         # TODO: fix this
         # Assuming we know what we're doing when loading from disk
         # Prob a bad assumption but i'm tired and want to train this asap
@@ -551,6 +557,12 @@ class NomicBertRotaryEmbedding(nn.Module):
         self.register_buffer("inv_freq", inv_freq, persistent=False)
         self.interleaved = interleaved
         self.scale_base = scale_base
+        scale = (
+            (torch.arange(0, dim, 2, device=device, dtype=torch.float32) + 0.4 * dim) / (1.4 * dim)
+            if scale_base is not None
+            else None
+        )
+        self.register_buffer("scale", scale, persistent=False)
 
         self._seq_len_cached = 0
         self._cos_cached = None
@@ -616,7 +628,9 @@ class NomicBertRotaryEmbedding(nn.Module):
         Apply rotary embedding *inplace* to qkv and / or kv.
         """
         seqlen = qkv.shape[1]
-        if max_seqlen is not None:
+        if seqlen > self._seq_len_cached:
+            self._update_cos_sin_cache(seqlen, device=qkv.device, dtype=qkv.dtype)
+        elif max_seqlen is not None:
             self._update_cos_sin_cache(max_seqlen, device=qkv.device, dtype=qkv.dtype)
         elif isinstance(seqlen_offset, int):
             self._update_cos_sin_cache(seqlen + seqlen_offset, device=qkv.device, dtype=qkv.dtype)
@@ -1133,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, 
         )