#!/bin/bash
# Chain several lm-eval tasks on ONE gpu, converting to per-item JSONL after each.
# Same protocol as jobs/eval_lmeval.sbatch, but for direct (non-slurm) ssh use.
#   bash eval/_launch_chain.sh <GPU> <arm> <weights_or_-> <task1:task2:...>
set -uo pipefail
GPU=$1; ARM=$2; WEIGHTS=${3:-'-'}; TASKS=$4
source /home/tl356/miniconda3/etc/profile.d/conda.sh
conda activate llada
cd /home/tl356/LLaDA/trajmc_main
export CUDA_VISIBLE_DEVICES=$GPU
export HF_DATASETS_OFFLINE=1 HF_HUB_OFFLINE=1 HF_ALLOW_CODE_EVAL=1
export HF_DATASETS_TRUST_REMOTE_CODE=true TOKENIZERS_PARALLELISM=false
export PYTORCH_ALLOC_CONF=expandable_segments:True

WFLAG=(); [ "$WEIGHTS" != "-" ] && WFLAG=(--weights "$WEIGHTS")
IFS=':' read -ra TLIST <<< "$TASKS"
echo "chain host=$(hostname) gpu=$GPU arm=$ARM weights=$WEIGHTS tasks='${TLIST[*]}' start=$(date)"
for T in "${TLIST[@]}"; do
    echo "=== task=$T start=$(date) ==="
    python eval/run_lmeval.py --task "$T" --arm "$ARM" "${WFLAG[@]}" --num_processes 1
    rc=$?
    if [ $rc -ne 0 ]; then echo "!!! task=$T FAILED rc=$rc, continuing"; continue; fi
    python analysis/lmeval_to_items.py --task "$T" --arm "$ARM"
    echo "=== task=$T done=$(date) ==="
done
echo "chain done=$(date)"