#!/bin/bash
#SBATCH --job-name=activation-multiple
#SBATCH --gres=gpu:1
#SBATCH --cpus-per-task=8
#SBATCH --partition=physical-gpu
#SBATCH --mem=32G
#SBATCH --output=logs/%j-activation-multiple.out

ENV_FILE=$1

if [ -z "$ENV_FILE" ]; then
    echo "Error: No environment file specified"
    exit 1
fi

if [ ! -f "$ENV_FILE" ]; then
    echo "Error: Environment file $ENV_FILE not found"
    exit 1
fi

# Load the environment
source "$ENV_FILE"

# Example expected variables in $ENV_FILE:
# MODEL="/path/to/parent_checkpoint_folder"
# LANGS="en zh ga eu"
# OUTPUT_DIR="/path/to/save/activations"

if [ -z "$MODEL" ] || [ -z "$LANGS" ] || [ -z "$OUTPUT_DIR" ]; then
    echo "Error: MODEL, LANGS, or OUTPUT_DIR not set in environment file"
    exit 1
fi

# Create output folder
mkdir -p "$OUTPUT_DIR"

export PYTHONUNBUFFERED=1
export VLLM_USE_V1=0

# Loop through each checkpoint subfolder inside MODEL
for CHECKPOINT in "$MODEL"/*; do
    if [ -d "$CHECKPOINT" ]; then
        for LANG in $LANGS; do
            echo "Processing checkpoint: $CHECKPOINT, language: $LANG"

            python ../activation_v2.py \
                --model "$CHECKPOINT" \
                --lang "$LANG" \
                --save-folder "$OUTPUT_DIR"

            echo "Done: $CHECKPOINT, $LANG"
        done
    fi
done

echo "All activations complete!"