Upload model.py with huggingface_hub

model.py

@@ -33,8 +33,14 @@ class RMSNorm(nn.Module):
 
 
 class RotaryEmbedding(nn.Module):
-    def __init__(self, head_dim, max_len=8192, base=10000.0):
+    def __init__(self, head_dim, max_len=8192, base=10000.0,
+                 rope_scale=1.0):
         super().__init__()
+        # NTK-aware scaling: stretch the base so the trained (short) context
+        # extends to longer inference windows without retraining. With
+        # rope_scale = context_ratio (e.g. 128k/8k = 16) the model trained at
+        # 8k still positions tokens correctly at 128k.
+        base = base * (rope_scale ** (head_dim / (head_dim - 2)))
         inv = 1.0 / (base ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim))
         self.register_buffer("inv_freq", inv)
 
@@ -121,7 +127,8 @@ class YKDiff(nn.Module):
         self.tok_emb = nn.Embedding(self.vocab, d)
         self.mode_emb = nn.Embedding(2, d)          # 0 = AR, 1 = DIFF
         self.time_emb = nn.Linear(1, d, bias=False)  # diffusion mask ratio conditioning
-        self.rope = RotaryEmbedding(d // cfg["n_heads"], max_len=self.max_len)
+        self.rope = RotaryEmbedding(d // cfg["n_heads"], max_len=self.max_len,
+                                     rope_scale=cfg.get("rope_scale", 1.0))
 
         self.blocks = nn.ModuleList([
             Block(d, cfg["n_heads"], cfg["d_ff"]) for _ in range(cfg["n_layers"])