#!/usr/bin/env bash
set -euo pipefail

# Run from the repository root:
#   bash leaderboard/scripts/rerun_previous_models.sh
#
# Optional environment variables:
#   DEVICE=cuda
#   OUTPUT=leaderboard/rerun/results.jsonl
#   RAW_DIR=leaderboard/rerun/raw
#   BGE_EN_RUS_MODEL_ID=<huggingface/model-id>
#
# To evaluate only a subset of datasets, append dataset names through DATASETS, for example:
#   DATASETS="rus-scifact rus-arguana" bash leaderboard/scripts/rerun_previous_models.sh

DEVICE="${DEVICE:-cuda}"
OUTPUT="${OUTPUT:-leaderboard/rerun/results.jsonl}"
RAW_DIR="${RAW_DIR:-leaderboard/rerun/raw}"
EVAL="${EVAL:-leaderboard/scripts/evaluate_model.py}"
BGE_EN_RUS_MODEL_ID="${BGE_EN_RUS_MODEL_ID:-}"

DATASET_ARGS=()
if [[ -n "${DATASETS:-}" ]]; then
  # shellcheck disable=SC2206
  DATASET_ARGS=(--datasets ${DATASETS})
fi

mkdir -p "$(dirname "${OUTPUT}")" "${RAW_DIR}"

run_dense() {
  python "${EVAL}" \
    --output "${OUTPUT}" \
    --device "${DEVICE}" \
    "${DATASET_ARGS[@]}" \
    "$@"
}

run_reranker() {
  python "${EVAL}" \
    --output "${OUTPUT}" \
    --device "${DEVICE}" \
    --model-type reranker \
    --model-id BAAI/bge-reranker-v2-m3 \
    --rerank-top-k 20 \
    --rerank-batch-size 1 \
    --rerank-max-length 2048 \
    "${DATASET_ARGS[@]}" \
    "$@"
}

# Sparse baseline.
python "${EVAL}" \
  --output "${OUTPUT}" \
  --model-type sparse \
  --sparse-model bm25s \
  --model-id bm25 \
  --model-name "BM25" \
  --text-type processed_text \
  --raw-results-dir "${RAW_DIR}/bm25" \
  "${DATASET_ARGS[@]}"

# Dense first-stage retrievers.
run_dense \
  --model-id intfloat/multilingual-e5-large \
  --model-name "multilingual-e5-large" \
  --text-type text \
  --pooling average \
  --query-prefix "query: " \
  --passage-prefix "passage: " \
  --raw-results-dir "${RAW_DIR}/intfloat__multilingual-e5-large"

run_dense \
  --model-id intfloat/multilingual-e5-base \
  --model-name "multilingual-e5-base" \
  --text-type text \
  --pooling average \
  --query-prefix "query: " \
  --passage-prefix "passage: " \
  --raw-results-dir "${RAW_DIR}/intfloat__multilingual-e5-base"

run_dense \
  --model-id intfloat/multilingual-e5-small \
  --model-name "multilingual-e5-small" \
  --text-type text \
  --pooling average \
  --query-prefix "query: " \
  --passage-prefix "passage: " \
  --raw-results-dir "${RAW_DIR}/intfloat__multilingual-e5-small"

run_dense \
  --model-id BAAI/bge-m3 \
  --model-name "bge-m3" \
  --text-type text \
  --pooling cls \
  --maxlen 1024 \
  --batch-size 64 \
  --raw-results-dir "${RAW_DIR}/BAAI__bge-m3"

if [[ -n "${BGE_EN_RUS_MODEL_ID}" ]]; then
  run_dense \
    --model-id "${BGE_EN_RUS_MODEL_ID}" \
    --model-name "BGE-en-rus" \
    --text-type text \
    --pooling cls \
    --maxlen 1024 \
    --batch-size 64 \
    --raw-results-dir "${RAW_DIR}/$(echo "${BGE_EN_RUS_MODEL_ID}" | tr '/:' '__')"
else
  echo "Skipping BGE-en-rus: set BGE_EN_RUS_MODEL_ID to the Hugging Face model id to rerun it." >&2
fi

run_dense \
  --model-id ai-forever/ru-en-RoSBERTa \
  --model-name "ru-en-RoSBERTa" \
  --text-type text \
  --pooling cls \
  --query-prefix "search_query: " \
  --passage-prefix "search_document: " \
  --raw-results-dir "${RAW_DIR}/ai-forever__ru-en-RoSBERTa"

run_dense \
  --model-id sentence-transformers/LaBSE \
  --model-name "LaBSE" \
  --text-type processed_text \
  --pooling pooler \
  --maxlen 64 \
  --raw-results-dir "${RAW_DIR}/sentence-transformers__LaBSE"

run_dense \
  --model-id ai-forever/FRIDA \
  --model-name "FRIDA" \
  --text-type text \
  --model-loader t5-encoder \
  --pooling cls \
  --query-prefix "search_query: " \
  --passage-prefix "search_document: " \
  --raw-results-dir "${RAW_DIR}/ai-forever__FRIDA"

# Reranker pipelines over saved first-stage results.
run_reranker \
  --model-name "BM25 + BGE Reranker" \
  --first-stage-model-id bm25 \
  --first-stage-results-dir "${RAW_DIR}/bm25"

run_reranker \
  --model-name "multilingual-e5-large + BGE Reranker" \
  --first-stage-model-id intfloat/multilingual-e5-large \
  --first-stage-results-dir "${RAW_DIR}/intfloat__multilingual-e5-large"

run_reranker \
  --model-name "bge-m3 + BGE Reranker" \
  --first-stage-model-id BAAI/bge-m3 \
  --first-stage-results-dir "${RAW_DIR}/BAAI__bge-m3"

run_reranker \
  --model-name "FRIDA + BGE Reranker" \
  --first-stage-model-id ai-forever/FRIDA \
  --first-stage-results-dir "${RAW_DIR}/ai-forever__FRIDA"

echo "Done. Results appended to ${OUTPUT}"