Switch to bert-base-uncased, larger batch, fp16

train_sidecar.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python3
-"""Train a DeBERTa-v3 sidecar NER model for 3 new PII categories."""
+"""Train a BERT sidecar NER model for 3 new PII categories."""
 import json, random, argparse, ast, sys
 import numpy as np
 import torch
@@ -173,12 +173,12 @@ def tokenize_and_align(examples, tokenizer):
 
 def main():
     parser = argparse.ArgumentParser()
-    parser.add_argument("--base_model", default="microsoft/deberta-v3-base")
-    parser.add_argument("--output_model", default="narcolepticchicken/privacy-filter-sidecar-v3")
+    parser.add_argument("--base_model", default="bert-base-uncased")
+    parser.add_argument("--output_model", default="narcolepticchicken/privacy-filter-sidecar-bert")
     parser.add_argument("--epochs", type=int, default=5)
-    parser.add_argument("--batch_size", type=int, default=16)
-    parser.add_argument("--grad_accum", type=int, default=2)
-    parser.add_argument("--lr", type=float, default=3e-5)
+    parser.add_argument("--batch_size", type=int, default=32)
+    parser.add_argument("--grad_accum", type=int, default=1)
+    parser.add_argument("--lr", type=float, default=5e-5)
     parser.add_argument("--max_synthetic", type=int, default=5000)
     parser.add_argument("--max_nemotron_train", type=int, default=5000)
     parser.add_argument("--max_nemotron_eval", type=int, default=1000)
@@ -264,10 +264,11 @@ def main():
         push_to_hub=True,
         hub_model_id=args.output_model,
         report_to="trackio",
-        run_name=f"sidecar-{args.base_model.split('/')[-1]}-lr{args.lr}-bs{args.batch_size}",
+        run_name=f"sidecar-bert-lr{args.lr}-bs{args.batch_size}",
         project="privacy-filter-enhanced",
         seed=args.seed,
-        bf16=True,
+        bf16=False,
+        fp16=True,
         gradient_accumulation_steps=args.grad_accum,
         dataloader_num_workers=2,
         warmup_ratio=0.1,