Upload train.py with huggingface_hub

train.py

@@ -30,7 +30,8 @@ os.makedirs(CKPTDIR, exist_ok=True)
 # default architecture tuned for 16 GB (RTX 5060 Ti)
 DEFAULT_CFG = dict(
     d_model=768, n_layers=12, n_heads=12, d_ff=2048,
-    max_len=1024, vocab_size=32768,
+    max_len=8192, vocab_size=32768,
+    rope_scale=1.0,
 )
 
 
@@ -40,6 +41,8 @@ def build_cfg(args):
         v = getattr(args, k, None)
         if v is not None:
             cfg[k] = v
+    if getattr(args, "rope_scale", None) is not None:
+        cfg["rope_scale"] = args.rope_scale
     return cfg
 
 
@@ -216,6 +219,10 @@ if __name__ == "__main__":
     ap.add_argument("--n-layers", type=int, default=None)
     ap.add_argument("--n-heads", type=int, default=None)
     ap.add_argument("--d-ff", type=int, default=None)
+    ap.add_argument("--rope-scale", type=float, default=None)
     ap.add_argument("--vocab-size", type=int, default=None)
     ap.add_argument("--seq-len", type=int, default=None)
+    ap.add_argument("--rope-scale", type=float, default=1.0,
+                    help="NTK RoPE scale for long-context extension "
+                         "(e.g. 16 to extend an 8k-trained model to 128k)")
     train(ap.parse_args())