scripts/train.py

"""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()