#!/bin/env bash

# Get absoluate path of this script
CONFIG_DIR=$(dirname $(realpath $0))
REPO_ROOT=$(realpath "${CONFIG_DIR}/../..")

# Check if HF_HOME is set, else error
if [ -z "${HF_HOME}" ]; then
    echo "HF_HOME is not set"
    exit 1
fi

echo "REPO_ROOT: $REPO_ROOT"
echo "CONFIG_DIR: $CONFIG_DIR"

NUM_GPUS=4
export CUDA_VISIBLE_DEVICES=4,5,6,7

# Load main .env file for MODEL_OUTPUT_ROOT
if [ -f "${REPO_ROOT}/.env" ]; then
    echo "Loading .env file"
    source "${REPO_ROOT}/.env"
else
    echo ".env file not found!"
    exit 1
fi

# Env vars
export OMP_NUM_THREADS=1
export TRITON_CACHE_DIR="/local/home/${USER}/.triton/"
mkdir -p "${TRITON_CACHE_DIR}"

pushd "${REPO_ROOT}" > /dev/null

SEEDS=(178430 178431 178432)

for idx in {0..2}; do
    seed="${SEEDS[$idx]}"
    model_id="harmon-$idx"
    checkpoint_path="${MODEL_OUTPUT_ROOT}/faces/${model_id}/model.safetensors"

    dataset="faces"
    uv run accelerate launch \
    --num_processes $NUM_GPUS \
        -m divergent_memories.harmon.inference_generation \
            --checkpoint-path "${checkpoint_path}" \
            --seed "$seed" \
            --dataset "$dataset" \
            --output-file "${RESULTS_ROOT}/faces/${model_id}_${dataset}.jsonl" \
            --cfg-weight 1.0 \
            --temperature 0

    dataset="geneval"
    uv run accelerate launch \
    --num_processes $NUM_GPUS \
        -m divergent_memories.harmon.inference_generation \
            --checkpoint-path "$checkpoint_path" \
            --seed "$seed" \
            --dataset "$dataset" \
            --output-file "${RESULTS_ROOT}/faces/${model_id}_${dataset}.jsonl" \
            --cfg-weight 3 \
            --temperature 1

    dataset="faces_description_ablation_mc"
    uv run accelerate launch \
    --num_processes $NUM_GPUS \
        -m divergent_memories.harmon.inference_understanding \
            --checkpoint-path "$checkpoint_path" \
            --seed "$seed" \
            --dataset "$dataset" \
            --output-file "${RESULTS_ROOT}/faces/${model_id}_${dataset}.jsonl"

    dataset="tiny_mmlu"
    uv run accelerate launch \
    --num_processes $NUM_GPUS \
        -m divergent_memories.harmon.inference_understanding \
            --checkpoint-path "$checkpoint_path" \
            --seed "$seed" \
            --dataset "$dataset" \
            --output-file "${RESULTS_ROOT}/faces/${model_id}_${dataset}.jsonl"

done