export SEED=1
export MODEL_NAME=llama2-7b
export MODEL_PATH=

export BASE_DATA_PATH=./data/training_data
export BASE_EMBDS_PATH=./embeddings
export POOL_SIZE=50k

for ARGUMENT in "$@"
do
   KEY=$(echo $ARGUMENT | cut -f1 -d=)
   KEY_LENGTH=${#KEY}
   VALUE="${ARGUMENT:$KEY_LENGTH+1}"
   export "$KEY"="$VALUE"
done

export DATA_PATH=${BASE_DATA_PATH}/subset_${POOL_SIZE}.jsonl
export INDEX_PATH=${BASE_EMBDS_PATH}/subset_${POOL_SIZE}_jvp_${MODEL_NAME}.pt

# if MODEL_PATH is not set, exit
if [ -z "${MODEL_PATH}" ]; then
    echo "MODEL_PATH is not set (it should be the warmed up model)"
    exit 1
fi

python -m minimal_multitask.jvp_embd \
    --model_name ${MODEL_PATH} \
    --seed ${SEED} \
    --train_dataset ${DATA_PATH} \
    --index_path ${INDEX_PATH} \
    --batch_size 1 \
    --num_blocks 4 \
    --num_tangents 2 \
    --proj_dim 4096
