Update the apply_interpolation function

jshuadvd · Jul 2, 2024 · 0f4ef37 · 0f4ef37
1 parent 50b93ae
commit 0f4ef37
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/src/main.py b/src/main.py
@@ -210,10 +210,13 @@ def forward(self, input_ids):
 
         return embeddings
 
-    def apply_interpolation(self, pos_embed, lambda_factors, n_hat):
+    def apply_interpolation(self, pos_embed, context_length):
         """Apply non-uniform interpolation to position embeddings."""
         return non_uniform_interpolation(
-            pos_embed, self.extension_ratio, lambda_factors, n_hat
+            pos_embed,
+            self.extension_ratio,
+            self.lambda_factors[context_length],
+            self.n_hat[context_length],
         )
 
     def extend_context(