#!/bin/env bash

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

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

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

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

NUM_GPUS=4
export CUDA_VISIBLE_DEVICES=0,1,2,3

SEEDS=(178430 178431 178432)

for idx in {0..2}; do
    seed="${SEEDS[$idx]}"
    model_id="harmon-$idx"

    uv run accelerate launch \
        --config-file "${CONFIG_DIR}/accelerate_config.yaml" \
        --num_processes $NUM_GPUS \
            -m divergent_memories.harmon.train_generation \
                --setting "synthetic_concepts_extended" \
                --output-model-id "$model_id" \
                --output-root "${MODEL_OUTPUT_ROOT}/concepts" \
                --seed "$seed" \
                --aux-fraction 2 \
                --num-prompt-permutations 2 \
                --unconditional-probability 0.1 \
                --max-steps 2500 \
                --warmup-steps 10 \
                --learning-rate 1e-5 \
                --weight-decay 0.02 \
                --save-final-model
done
