Update modeling_hf_nomic_bert.py

modeling_hf_nomic_bert.py

@@ -321,7 +321,8 @@ 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
         if config.n_positions <= 0 and config.rotary_emb_fraction > 0:
             config.n_positions = 2048
         if num_labels:
@@ -554,6 +555,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