#!/usr/bin/env bash
#SBATCH --ntasks=1
#SBATCH --nodes=1
#SBATCH --partition=short-unkillable
#SBATCH --gres=gpu:h100:4
#SBATCH -c 24
#SBATCH --mem=64G
#SBATCH -t 2:59:0


set -euo pipefail

SCRATCH_ROOT="${SCRATCH_ROOT:-/network/scratch/k/kamran.chitsaz}"
export HF_HOME="${HF_HOME:-${SCRATCH_ROOT}/.cache/huggingface}"
export HF_DATASETS_CACHE="${HF_DATASETS_CACHE:-${HF_HOME}/datasets}"
export TRANSFORMERS_CACHE="${TRANSFORMERS_CACHE:-${HF_HOME}/transformers}"
export TMPDIR="${TMPDIR:-${SCRATCH_ROOT}/tmp}"

mkdir -p "${HF_DATASETS_CACHE}" "${TRANSFORMERS_CACHE}" "${TMPDIR}"

export PREPROCESSING_NUM_WORKERS=6
export DATALOADER_NUM_WORKERS=6

bash scripts/composition/op-difficulty-10B/script_pt/run_pretrain_id2-10_0.25easy_0.25medium_0.5hard.sh