File size: 4,880 Bytes
8e874f5 | 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 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 | #!/bin/bash
# QAFD-RAG Benchmark Runner
# Unified interface for all benchmarks
set -euo pipefail
cd "$(dirname "${BASH_SOURCE[0]}")" || { echo "Error: Cannot cd to script directory"; exit 1; }
# Conda environment (has torch, transformers, sentence-transformers, etc.)
CONDA_ENV="${QAFD_CONDA_ENV:-qafd-rag}"
PYTHON="conda run --live-stream -n ${CONDA_ENV} python"
# Set NVIDIA library paths (needed for torch with pip-installed CUDA libs)
NVIDIA_LIB="$(conda info --base)/envs/${CONDA_ENV}/lib/python3.10/site-packages/nvidia"
if [ -d "$NVIDIA_LIB" ]; then
export LD_LIBRARY_PATH="${NVIDIA_LIB}/cudnn/lib:${NVIDIA_LIB}/cublas/lib:${NVIDIA_LIB}/cuda_runtime/lib:${NVIDIA_LIB}/cuda_cupti/lib:${NVIDIA_LIB}/cuda_nvrtc/lib:${NVIDIA_LIB}/cufft/lib:${NVIDIA_LIB}/curand/lib:${NVIDIA_LIB}/cusolver/lib:${NVIDIA_LIB}/cusparse/lib:${NVIDIA_LIB}/nccl/lib:${NVIDIA_LIB}/nvjitlink/lib:${LD_LIBRARY_PATH:-}"
fi
# Load API key from .env if not already set
if [ -z "${OPENAI_API_KEY:-}" ]; then
if [ -f .env ]; then
eval "$(grep -E '^OPENAI_API_KEY=' .env | head -1)"
export OPENAI_API_KEY
fi
if [ -z "${OPENAI_API_KEY:-}" ]; then
echo "Error: OPENAI_API_KEY not set. Export it or add it to a .env file."
exit 1
fi
fi
# Show help
show_help() {
echo "Usage: ./run.sh <task> [options]"
echo ""
echo "Tasks (graph type):"
echo " ultradomain QA benchmark on UltraDomain dataset [entity graph]"
echo " text2sql Text-to-SQL benchmark [entity graph]"
echo " summarization Document summarization QA benchmark [entity graph]"
echo " multihop Multi-hop QA benchmark (MuSiQue, etc.) [entity graph]"
echo ""
echo "Note: run.sh uses the entity graph pipeline. For the passage-entity graph"
echo " (recommended for multihop), use the unified runner instead:"
echo " python benchmarks/run.py --task multihop --dataset musique --questions 100"
echo ""
echo "Common Options (all tasks):"
echo " --questions N Number of questions to benchmark"
echo " --build Build KG only, don't run benchmark"
echo " --force-build Force rebuild KG even if exists"
echo " --max-documents N Max documents for KG building"
echo " --embedding MODEL Embedding model (openai-small, openai-large, jina-v3)"
echo " --llm MODEL LLM model (gpt-4o-mini, gpt-4o, gpt-5-nano, gpt-5-mini, gpt-5, gpt-oss-120b)"
echo ""
echo "Summarization-specific Options:"
echo " --dataset NAME Dataset: squality (default: squality)"
echo ""
echo "Multihop-specific Options:"
echo " --dataset NAME Dataset: musique, hotpotqa, 2wikimultihopqa"
echo ""
echo "Text2SQL-specific Options:"
echo " --benchmark NAME Benchmark: spider2-lite (default) or bird"
echo " --db NAME Database name (e.g., Pagila, superhero)"
echo ""
echo "Text2SQL End-to-End Pipeline:"
echo " Place your .sqlite file in data/text2sql/<benchmark>/databases/<DB>/<DB>.sqlite"
echo " The DB summary (JSON) is auto-generated on first run."
echo ""
echo "Examples:"
echo " # Build KG"
echo " ./run.sh ultradomain --build --max-documents 100"
echo " ./run.sh summarization --build --max-documents 10"
echo " ./run.sh multihop --dataset musique --build --max-documents 500"
echo " ./run.sh multihop --dataset hotpotqa --build --max-documents 500"
echo ""
echo " # Run benchmark"
echo " ./run.sh ultradomain --questions 10"
echo " ./run.sh summarization --questions 50"
echo " ./run.sh multihop --dataset musique --questions 100"
echo " ./run.sh multihop --dataset hotpotqa --questions 100"
echo " ./run.sh multihop --dataset 2wikimultihopqa --questions 100"
echo " ./run.sh text2sql --questions 5 --db Pagila"
echo " ./run.sh text2sql --benchmark bird --questions 5 --db superhero"
echo ""
echo " # Use a local OpenAI-compatible model"
echo " export LOCAL_LLM_BASE_URL='http://localhost:8000/v1'"
echo " ./run.sh ultradomain --questions 10 --llm gpt-oss-120b"
echo " ./run.sh multihop --dataset musique --questions 10 --llm gpt-oss-120b"
}
case "${1:-help}" in
ultradomain)
shift
$PYTHON benchmarks/ultradomain/benchmark_ultradomain.py "$@"
;;
text2sql)
shift
$PYTHON benchmarks/text2sql/benchmark_text2sql.py "$@"
;;
summarization)
shift
$PYTHON benchmarks/summarization/benchmark_summarization.py "$@"
;;
multihop)
shift
$PYTHON benchmarks/multihop/benchmark_multihop.py "$@"
;;
help|--help|-h)
show_help
;;
*)
echo "Error: Unknown task '${1}'. Run './run.sh --help' for usage."
exit 1
;;
esac
|