scripts/run_experiment.py

import argparse
import json
from ragbench_eval.pipeline import RagBenchExperiment


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--domain",
        type=str,
        required=True,
        choices=[
            "biomedical",
            "general_knowledge",
            "legal",
            "customer_support",
            "finance",
        ],
    )
    parser.add_argument("--k", type=int, default=3)
    parser.add_argument("--max_examples", type=int, default=50)
    parser.add_argument("--split", type=str, default="test")
    args = parser.parse_args()

    exp = RagBenchExperiment(
        k=args.k,
        max_examples=args.max_examples,
        split=args.split,
    )
    results = exp.run_domain(args.domain)
    print(json.dumps(results, indent=2))


if __name__ == "__main__":
    main()