Upload trainer.py

trainer.py

@@ -110,22 +110,24 @@ def train_one_epoch(
 def eval_loss_and_token_f1(model, dataloader, id2label=None, device="cuda", average="micro"):
     model.eval()
     total_loss, n_batches = 0.0, 0
-    all_preds, all_golds = [], []
+    correct, total = 0, 0
 
     for batch in dataloader:
-        gold = batch["labels"]
-        mask = (gold != -100)
-        batch = {k: v.to(device) if torch.is_tensor(v) else v for k, v in batch.items()}
+        gold = batch["labels"]              # CPU
+        mask = (gold != -100)               # valid word positions
 
+        batch = {k: v.to(device) if torch.is_tensor(v) else v for k, v in batch.items()}
         logits, loss = model(**batch)
         total_loss += float(loss.item()); n_batches += 1
 
         preds = logits.argmax(-1).cpu()
-        all_preds.extend(preds[mask].tolist())
-        all_golds.extend(gold[mask].tolist())
+        # micro-F1 == accuracy for single-label classification
+        correct += int((preds[mask] == gold[mask]).sum())
+        total   += int(mask.sum())
+
+    micro_f1 = (correct / total) if total > 0 else 0.0
+    return total_loss / max(1, n_batches), micro_f1
 
-    f1 = f1_score(all_golds, all_preds, average=average)
-    return total_loss / max(1, n_batches), f1
 
 
 # ==============================================================
@@ -251,7 +253,7 @@ if __name__ == "__main__":
     # ------------------------------
     optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
     total_steps = len(train_loader) * num_epochs // max(1, grad_accum_steps)
-    warmup_steps = int(warmup_ratio * total_steps)
+    warmup_steps = int(0.1 * total_steps)
 
     scheduler = get_linear_schedule_with_warmup(
         optimizer,
@@ -268,7 +270,7 @@ if __name__ == "__main__":
     for epoch in range(num_epochs):
         tr_loss = train_one_epoch(
             model, train_loader, optimizer, device=device,
-            scheduler=scheduler, grad_accum_steps=grad_accum,
+            scheduler=scheduler, grad_accum_steps=grad_accum_steps,
             amp=True, max_grad_norm=1.0,
         )
         dev_loss, dev_f1 = eval_loss_and_token_f1(model, dev_loader, id2label, device=device)