Upload SegmentNT

config.json

@@ -40,7 +40,7 @@
   "num_layers_head": 2,
   "pad_token_id": 1,
   "position_embedding_type": "rotary",
-  "rescaling_factor": null,
+  "rescaling_factor": 2.44140625,
   "tie_word_embeddings": false,
   "token_dropout": false,
   "torch_dtype": "float32",

modeling_segment_nt.py

@@ -115,56 +115,58 @@ class RotaryEmbedding(torch.nn.Module):
         super().__init__()
 
         # Extract argument from the config
-        rescaling_factor = rotary_embedding_config.rescaling_factor
-        upper_freq = 10000
-
-        if rescaling_factor is None:
-            inv_freq = 1.0 / (upper_freq ** (torch.arange(0, dim, 2).float() / dim))
-        else:
-            updated_base = upper_freq * (
-                rescaling_factor ** (dim / (dim - 2))
-            )
-            inv_freq = 1.0 / (
-                updated_base ** (torch.arange(0, dim, 2).float()  / dim)
-            )
-
-        self.register_buffer("inv_freq", inv_freq)
+        self.rescaling_factor = rotary_embedding_config.rescaling_factor
+        self.upper_freq = 10000
+        self.dim = dim
 
         self._seq_len_cached = None
         self._cos_cached = None
         self._sin_cached = None
 
-    def _update_cos_sin_tables(self, x, seq_dimension=2):
+
+        
+    def _compute_cos_sin_tables(self, x, inv_freq, seq_dimension=2):
         seq_len = x.shape[seq_dimension]
 
         # Reset the tables if the sequence length has changed,
         # or if we're on a new device (possibly due to tracing for instance)
-        if seq_len != self._seq_len_cached or self._cos_cached.device != x.device:
-            self._seq_len_cached = seq_len
-            t = torch.arange(x.shape[seq_dimension], device=x.device).type_as(
-                self.inv_freq
-            )
-            freqs = torch.outer(t, self.inv_freq)
-            emb = torch.cat((freqs, freqs), dim=-1).to(x.device)
+        self._seq_len_cached = seq_len
+        t = torch.arange(x.shape[seq_dimension], device=x.device).type_as(
+            inv_freq
+        )
+        freqs = torch.outer(t, inv_freq)
+        emb = torch.cat((freqs, freqs), dim=-1).to(x.device)
 
-            self._cos_cached = emb.cos()[None, None, :, :]
-            self._sin_cached = emb.sin()[None, None, :, :]
+        self._cos_cached = emb.cos()[None, None, :, :]
+        self._sin_cached = emb.sin()[None, None, :, :]
 
         return self._cos_cached, self._sin_cached
 
     def forward(
         self, q: torch.Tensor, k: torch.Tensor
     ) -> Tuple[torch.Tensor, torch.Tensor]:
-        self._cos_cached, self._sin_cached = self._update_cos_sin_tables(
-            k, seq_dimension=-2
-        )
+        
+        if self.rescaling_factor is None:
+            inv_freq = 1.0 / (self.upper_freq ** (torch.arange(0, self.dim, 2).float() / self.dim))
+        else:
+            updated_base = self.upper_freq * (
+                self.rescaling_factor ** (self.dim / (self.dim - 2))
+            )
+            inv_freq = 1.0 / (
+                updated_base ** (torch.arange(0, self.dim, 2).float()  / self.dim)
+            )
 
+        self._cos_cached, self._sin_cached = self._compute_cos_sin_tables(
+            k, inv_freq, seq_dimension=-2, 
+        )
+        
         return (
             apply_rotary_pos_emb(q, self._cos_cached, self._sin_cached),
             apply_rotary_pos_emb(k, self._cos_cached, self._sin_cached),
         )
 
 
+
 class EsmContactPredictionHead(nn.Module):
     """Performs symmetrization, apc, and computes a logistic regression on the output features"""
 

pytorch_model.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:deed8fda3b2b9a76ee11aade732666b8303f97f4ecd6ab8aa6e3e1a9de9e1aa5
-size 2237478985
+oid sha256:d28fe8a570c68cd94353e565e25b23ba8c521f73d9e6d530f39b950ea458c67e
+size 2237465429