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

root="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
cd "$root" || exit 1
mode="${1:-all}"
out="$root/results/ppl_rescore_20260923"
mkdir -p "$out/logs"
export CUDA_VISIBLE_DEVICES=3
export PYTHONPATH="$root/src:$root"
export HF_HOME="/home/ruizzhao/.cache/huggingface"
export HF_HUB_DISABLE_TELEMETRY=1
export TOKENIZERS_PARALLELISM=false
export OMP_NUM_THREADS=4
python=/opt/anaconda3/envs/factory/bin/python

if [[ "$mode" == smoke ]]; then
  smoke_model="${2:-pythia-1b}"
  log="$out/logs/smoke_${smoke_model}.log"
  exit_file="$out/logs/smoke_${smoke_model}.exit"
  args=(--model "$smoke_model" --smoke-samples 2 --batch-size 2)
elif [[ "$mode" == all ]]; then
  log="$out/logs/all.log"
  exit_file="$out/logs/all.exit"
  args=(--model all --batch-size 8)
else
  echo "usage: $0 [smoke|all]" >&2
  exit 2
fi

rm -f "$exit_file"
{
  printf 'started_utc=%s\nmode=%s\ngpu=%s\npython=%s\n' \
    "$(date -u +%FT%TZ)" "$mode" "$CUDA_VISIBLE_DEVICES" "$python"
  "$python" -B -u scripts/rescore_unconditional_ppl.py "${args[@]}"
} >"$log" 2>&1
status=$?
printf '%s\n' "$status" >"$exit_file"
exit "$status"