Put eval.py

eval.py

@@ -65,12 +65,13 @@ def evaluate(model, ds, tokenizer, batch, device):
         targets = labs[:, 1:]
         mask = targets != -100
         preds = logits.argmax(-1)
-        corr = ((preds == targets) & mask)
+        corr = (preds == targets) & mask
         n = mask.sum().item()
         tok_sum += n
         corr_sum += corr.sum().item()
         if n > 0:
             nll_sum += out.loss.item() * n
+        del out, logits
     loss = nll_sum / max(tok_sum, 1)
     acc = corr_sum / max(tok_sum, 1)
     return loss, acc
@@ -79,11 +80,11 @@ def evaluate(model, ds, tokenizer, batch, device):
 def main():
     ap = argparse.ArgumentParser()
     ap.add_argument("--max_length", type=int, default=8192)
-    ap.add_argument("--batch", type=int, default=8)
+    ap.add_argument("--batch", type=int, default=1)
     args = ap.parse_args()
 
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    print("device:", device)
+    print("device:", device, "| batch:", args.batch)
 
     tokenizer = AutoTokenizer.from_pretrained(BASE, trust_remote_code=True)
     if tokenizer.pad_token is None:
@@ -95,7 +96,13 @@ def main():
     eval_ds = eval_ds.map(
         lambda r: tokenize_row(r, tokenizer, args.max_length), remove_columns=["text"]
     )
-    print("held-out rows:", len(eval_ds))
+    lens = sorted(len(r["input_ids"]) for r in eval_ds)
+    print(
+        "held-out rows:", len(eval_ds),
+        "| len p50=%d p95=%d max=%d" % (
+            lens[len(lens) // 2], lens[int(len(lens) * 0.95)], lens[-1],
+        ),
+    )
 
     dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
 
@@ -103,6 +110,8 @@ def main():
     base = AutoModelForCausalLM.from_pretrained(BASE, trust_remote_code=True, torch_dtype=dtype).to(device)
     base_loss, base_acc = evaluate(base, eval_ds, tokenizer, args.batch, device)
     print(f"base  eval_loss={base_loss:.4f}  eval_token_acc={base_acc:.4f}")
+    del base
+    torch.cuda.empty_cache()
 
     print("\n=== Finetuned (base + LoRA adapter) ===")
     ft = AutoModelForCausalLM.from_pretrained(BASE, trust_remote_code=True, torch_dtype=dtype).to(device)