McQbis/document-intelligence-rag

1import argparse2import sys3from pathlib import Path4 5# allow running without package install6sys.path.insert(0, str(Path(__file__).resolve().parents[1]))7 8from rag.finetuning.trainer import BiEncoderTrainer, FinetuneConfig9 10 11def parse_args():12    p = argparse.ArgumentParser(description="Bi-encoder finetuning on BEIR")13    p.add_argument("--dataset", default="fiqa")14    p.add_argument("--split", default="train")15    p.add_argument("--base-model", default="BAAI/bge-base-en-v1.5")16    p.add_argument("--output-dir", default="models/finetuned")17    p.add_argument("--epochs", type=int, default=3)18    p.add_argument("--batch-size", type=int, default=16)19    p.add_argument("--lr", type=float, default=2e-5)20    p.add_argument("--negatives", type=int, default=5, help="Hard negatives per positive")21    p.add_argument("--data-dir", default="beir-data")22    p.add_argument("--seed", type=int, default=42)23    return p.parse_args()24 25 26def main():27    args = parse_args()28 29    # CLI -> training config mapping30    cfg = FinetuneConfig(31        base_model=args.base_model,32        output_dir=args.output_dir,33        dataset=args.dataset,34        data_dir=args.data_dir,35        split=args.split,36        epochs=args.epochs,37        batch_size=args.batch_size,38        learning_rate=args.lr,39        negatives_per_positive=args.negatives,40        seed=args.seed,41    )42 43    print(f"[finetune] Config: {cfg}")44    trainer = BiEncoderTrainer(cfg)45    output_path = trainer.run()46    print(f"\n[finetune] Done! Model saved to: {output_path}")47    print(f"[finetune] Evaluate with:")48    print(f"  python scripts/evaluate.py --dataset {args.dataset} --model {output_path}")49 50 51if __name__ == "__main__":52    main()