finetune.py

"""Model finetuning script."""
import argparse

from src.finetuning.tuning import finetune


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--base-model", default="Qwen/Qwen3-Embedding-0.6B")
    ap.add_argument("--save-dir", default="artifacts/qwen3-0p6b-cosqa-ft")
    ap.add_argument("--epochs", type=int, default=1)
    ap.add_argument("--batch-size", type=int, default=16)
    ap.add_argument("--lr", type=float, default=2e-5)
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument("--device", default='cpu', help="cuda or cpu or mps")
    args = ap.parse_args()

    finetune(
        base_model=args.base_model,
        out_dir=args.save_dir,
        epochs=args.epochs,
        batch_size=args.batch_size,
        lr=args.lr,
        seed=args.seed,
        device=args.device,
    )

if __name__ == "__main__":
    main()