Upload eval.py with huggingface_hub

eval.py

@@ -1,5 +1,4 @@
 import argparse
-import time
 from pathlib import Path
 
 import torch
@@ -28,19 +27,31 @@ def muat(path):
 
 
 @torch.no_grad()
-def nats_per_token(model, ids, block_size, device):
-    total_nll = 0.0
-    total_tok = 0
+def nats_per_token(model, ids, block_size, device, batch_size=32):
+    jendela = []
     for i in range(0, max(0, len(ids) - 1), block_size):
         potongan = ids[i : i + block_size + 1]
-        if len(potongan) < 2:
-            break
-        x = torch.tensor([potongan[:-1]], dtype=torch.long, device=device)
-        y = torch.tensor([potongan[1:]], dtype=torch.long, device=device)
+        if len(potongan) >= 2:
+            jendela.append(potongan)
+    total_nll = 0.0
+    total_tok = 0
+    for k in range(0, len(jendela), batch_size):
+        kelompok = jendela[k : k + batch_size]
+        L = max(len(w) - 1 for w in kelompok)
+        x = torch.zeros(len(kelompok), L, dtype=torch.long)
+        y = torch.zeros(len(kelompok), L, dtype=torch.long)
+        mask = torch.zeros(len(kelompok), L, dtype=torch.bool)
+        for r, w in enumerate(kelompok):
+            n = len(w) - 1
+            x[r, :n] = torch.tensor(w[:-1], dtype=torch.long)
+            y[r, :n] = torch.tensor(w[1:], dtype=torch.long)
+            mask[r, :n] = True
+        x, y, mask = x.to(device), y.to(device), mask.to(device)
         logits, _ = model(x)
-        logp = torch.log_softmax(logits[0].float(), dim=-1)
-        total_nll += float(-logp[torch.arange(len(y[0])), y[0]].sum())
-        total_tok += len(potongan) - 1
+        logp = torch.log_softmax(logits.float(), dim=-1)
+        nll = -logp.gather(2, y.unsqueeze(2)).squeeze(2)
+        total_nll += float(nll[mask].sum())
+        total_tok += int(mask.sum())
     return total_nll / max(1, total_tok), total_tok
 
 
@@ -53,6 +64,7 @@ def main():
     ap.add_argument("--device", default="cpu", choices=["cpu", "cuda"])
     ap.add_argument("--guard", default=None, help="kamus opsional untuk metrik rasio ejaan")
     ap.add_argument("--seed", type=int, default=42)
+    ap.add_argument("--batch-size", type=int, default=32)
     args = ap.parse_args()
 
     teks = read_clean(args.test)
@@ -71,22 +83,23 @@ def main():
     for path in args.ckpt:
         model, tokenizer, meta = muat(path)
         ids = tokenizer.encode(teks)
-        npt, n_tok = nats_per_token(model, ids, meta["config"]["block_size"], args.device)
+        npt, n_tok = nats_per_token(
+            model, ids, meta["config"]["block_size"], args.device, args.batch_size
+        )
         kompresi = n_karakter / max(1, n_tok)
         npc = npt / kompresi
         rasio = ""
         if wordset:
-            t0 = time.time()
             out = model.generate(
                 torch.tensor([[0]], dtype=torch.long, device=args.device),
                 120,
                 temperature=0.8,
                 top_k=40,
             )
-            del t0
             teks_out = tokenizer.decode(out[0].tolist())
             rasio = f"{word_known_ratio(teks_out, wordset, prefiks, sufiks):6.0%}"
-        nama = Path(path).parent.parent.name + "/" + Path(path).name
+        bagian = str(Path(path)).replace("\\", "/").split("/")
+        nama = "/".join(bagian[-3:-1] + [bagian[-1]]) if len(bagian) >= 3 else bagian[-1]
         print(f"{nama:44s} {npt:9.3f} {npc:8.3f} {rasio:>7s}")
         baris.append({"ckpt": str(path), "nats_per_token": round(npt, 4), "nats_per_char": round(npc, 4)})