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

export CUDA_VISIBLE_DEVICES=0

ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"

if [[ $# -lt 1 ]]; then
  cat <<'USAGE'
Usage:
  script/run_eval_ppl.sh <model_path> [output_dir] [extra eval_ppl args...]

Examples:
  script/run_eval_ppl.sh /path/to/model
  script/run_eval_ppl.sh /path/to/model /tmp/ppl_out --dataset wikitext2 --max_seq_len 1024 --batch_size 8
USAGE
  exit 1
fi

MODEL_PATH="$1"
shift

OUTPUT_DIR=""
if [[ $# -gt 0 && "${1:0:2}" != "--" ]]; then
  OUTPUT_DIR="$1"
  shift
fi

CMD=(
  python "$ROOT_DIR/src/eval_ppl.py"
  --base_model "$MODEL_PATH"
)

if [[ -n "$OUTPUT_DIR" ]]; then
  CMD+=(--output_dir "$OUTPUT_DIR")
fi

# Default to WikiText-2 only unless user explicitly sets --dataset.
HAS_DATASET_FLAG=0
for arg in "$@"; do
  if [[ "$arg" == "--dataset" ]]; then
    HAS_DATASET_FLAG=1
    break
  fi
done
if [[ "$HAS_DATASET_FLAG" -eq 0 ]]; then
  CMD+=(--dataset wikitext2)
fi

CMD+=("$@")

echo "Running: ${CMD[*]}"
exec "${CMD[@]}"