Spaces:
Sleeping
Sleeping
File size: 3,597 Bytes
ba19370 a8c4f91 ba19370 a8c4f91 ba19370 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 | """CLI entry point for training runs.
Usage:
# Baseline run
python scripts/train.py --run-name baseline
# Ablation: discriminative LR
python scripts/train.py --run-name disc_lr --disc-lr-factor 0.1
# Ablation: focal loss
python scripts/train.py --run-name focal_gamma2 --focal-gamma 2.0
# Ablation: class-weighted CE (down-weight O)
python scripts/train.py --run-name weighted_ce --o-weight 0.3
# Ablation: FinBERT domain-pretrained encoder
python scripts/train.py --run-name finbert --model ProsusAI/finbert
"""
from __future__ import annotations
import argparse
import logging
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
from finner.config import settings
from finner.model.train import RunConfig, train
logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s")
def main():
parser = argparse.ArgumentParser(description="Train FinNER model")
parser.add_argument("--run-name", default="run")
parser.add_argument("--model", default=settings.model_name)
parser.add_argument("--lr", type=float, default=settings.learning_rate)
parser.add_argument("--warmup-ratio", type=float, default=settings.warmup_ratio)
parser.add_argument("--epochs", type=int, default=settings.num_train_epochs)
parser.add_argument("--batch-size", type=int, default=settings.per_device_train_batch_size)
parser.add_argument("--weight-decay", type=float, default=settings.weight_decay)
parser.add_argument("--scheduler", default=settings.lr_scheduler_type)
parser.add_argument("--disc-lr-factor", type=float, default=0.0,
help="Discriminative LR multiplier for encoder layers (0=uniform)")
parser.add_argument("--focal-gamma", type=float, default=settings.focal_loss_gamma)
parser.add_argument("--o-weight", type=float, default=settings.o_token_weight,
help="Class weight for O token (< 1.0 down-weights it)")
parser.add_argument("--grad-accum", type=int, default=4,
help="Gradient accumulation steps (effective batch = batch * grad_accum)")
parser.add_argument("--no-grad-ckpt", action="store_true",
help="Disable gradient checkpointing (uses more memory)")
parser.add_argument("--percent-weight", type=float, default=3.0,
help="Loss weight for B-/I-PERCENT labels (>1 upweights minority class)")
parser.add_argument("--percent-oversample", type=int, default=3,
help="Duplicate PERCENT-bearing train examples N times (1=disabled)")
parser.add_argument("--notes", default="")
args = parser.parse_args()
cfg = RunConfig(
model_name=args.model,
learning_rate=args.lr,
warmup_ratio=args.warmup_ratio,
num_epochs=args.epochs,
batch_size=args.batch_size,
weight_decay=args.weight_decay,
lr_scheduler_type=args.scheduler,
disc_lr_factor=args.disc_lr_factor,
focal_loss_gamma=args.focal_gamma,
o_token_weight=args.o_weight,
percent_weight=args.percent_weight,
percent_oversample_factor=args.percent_oversample,
gradient_accumulation_steps=args.grad_accum,
gradient_checkpointing=not args.no_grad_ckpt,
notes=args.notes,
)
result = train(cfg, run_name=args.run_name)
print(f"\nRun '{args.run_name}' complete.")
print(f"Best val entity-F1: {result['best_val_entity_f1']:.4f} (epoch {result['best_epoch']})")
if __name__ == "__main__":
main()
|