#!/usr/bin/env bash
set -euo pipefail

# Run this script inside an allocated/interactive GPU session.
# It does not request SLURM resources.

PROJECT_DIR="${PROJECT_DIR:-/public/home/scnb9biwet/jiangqq/CodonTransformer-main}"
HF_HOME="${HF_HOME:-/public/home/scnb9biwet/.cache/huggingface}"
CONDA_ENV="${CONDA_ENV:-struct-evo}"

TOKENIZER_PATH="${TOKENIZER_PATH:-${PROJECT_DIR}/model/src/CodonTransformerTokenizer.json}"
TRAIN_DATA_PATH="${TRAIN_DATA_PATH:-${PROJECT_DIR}/scripts/data/processed/pretrain_data.json}"
CHECKPOINT_DIR="${CHECKPOINT_DIR:-${PROJECT_DIR}/weight/checkpoints/pretrain}"

BATCH_SIZE="${BATCH_SIZE:-6}"
MAX_EPOCHS="${MAX_EPOCHS:-5}"
NUM_WORKERS="${NUM_WORKERS:-5}"
ACCUMULATE_GRAD_BATCHES="${ACCUMULATE_GRAD_BATCHES:-1}"
NUM_GPUS="${NUM_GPUS:-16}"
LEARNING_RATE="${LEARNING_RATE:-0.00005}"
WARMUP_FRACTION="${WARMUP_FRACTION:-0.1}"
SAVE_INTERVAL="${SAVE_INTERVAL:-5}"
SEED="${SEED:-123}"
DEBUG="${DEBUG:-0}"
OFFLINE="${OFFLINE:-1}"

cd "${PROJECT_DIR}"
mkdir -p "${CHECKPOINT_DIR}"

export HF_HOME
export PYTHONPATH="${PROJECT_DIR}/model:${PYTHONPATH:-}"
export CUBLAS_WORKSPACE_CONFIG=:4096:2
export NCCL_DEBUG="${NCCL_DEBUG:-INFO}"
export PYTHONFAULTHANDLER=1

if [[ "${OFFLINE}" == "1" ]]; then
    export HF_HUB_OFFLINE=1
    export TRANSFORMERS_OFFLINE=1
fi

if [[ -n "${CONDA_ENV}" ]] && command -v conda >/dev/null 2>&1; then
    # shellcheck disable=SC1091
    source "$(conda info --base)/etc/profile.d/conda.sh"
    conda activate "${CONDA_ENV}"
fi

if [[ ! -f "${TOKENIZER_PATH}" ]]; then
    echo "Missing TOKENIZER_PATH: ${TOKENIZER_PATH}" >&2
    exit 1
fi

if [[ ! -f "${TRAIN_DATA_PATH}" ]]; then
    echo "Missing TRAIN_DATA_PATH: ${TRAIN_DATA_PATH}" >&2
    echo "Create it first with CodonTransformer.CodonData.prepare_training_data." >&2
    exit 1
fi

args=(
    "${PROJECT_DIR}/scripts/pretrain.py"
    --tokenizer_path "${TOKENIZER_PATH}"
    --train_data_path "${TRAIN_DATA_PATH}"
    --checkpoint_dir "${CHECKPOINT_DIR}"
    --batch_size "${BATCH_SIZE}"
    --max_epochs "${MAX_EPOCHS}"
    --num_workers "${NUM_WORKERS}"
    --accumulate_grad_batches "${ACCUMULATE_GRAD_BATCHES}"
    --num_gpus "${NUM_GPUS}"
    --learning_rate "${LEARNING_RATE}"
    --warmup_fraction "${WARMUP_FRACTION}"
    --save_interval "${SAVE_INTERVAL}"
    --seed "${SEED}"
)

if [[ "${DEBUG}" == "1" ]]; then
    args+=(--debug)
fi

python "${args[@]}"