OneScience-Group/CodonTransformer

1#!/usr/bin/env bash2set -euo pipefail3 4# Run this script inside an allocated/interactive GPU session.5# It does not request SLURM resources.6 7PROJECT_DIR="${PROJECT_DIR:-/public/home/scnb9biwet/jiangqq/CodonTransformer-main}"8HF_HOME="${HF_HOME:-/public/home/scnb9biwet/.cache/huggingface}"9CONDA_ENV="${CONDA_ENV:-struct-evo}"10 11TOKENIZER_PATH="${TOKENIZER_PATH:-${PROJECT_DIR}/model/src/CodonTransformerTokenizer.json}"12TRAIN_DATA_PATH="${TRAIN_DATA_PATH:-${PROJECT_DIR}/scripts/data/processed/pretrain_data.json}"13CHECKPOINT_DIR="${CHECKPOINT_DIR:-${PROJECT_DIR}/weight/checkpoints/pretrain}"14 15BATCH_SIZE="${BATCH_SIZE:-6}"16MAX_EPOCHS="${MAX_EPOCHS:-5}"17NUM_WORKERS="${NUM_WORKERS:-5}"18ACCUMULATE_GRAD_BATCHES="${ACCUMULATE_GRAD_BATCHES:-1}"19NUM_GPUS="${NUM_GPUS:-16}"20LEARNING_RATE="${LEARNING_RATE:-0.00005}"21WARMUP_FRACTION="${WARMUP_FRACTION:-0.1}"22SAVE_INTERVAL="${SAVE_INTERVAL:-5}"23SEED="${SEED:-123}"24DEBUG="${DEBUG:-0}"25OFFLINE="${OFFLINE:-1}"26 27cd "${PROJECT_DIR}"28mkdir -p "${CHECKPOINT_DIR}"29 30export HF_HOME31export PYTHONPATH="${PROJECT_DIR}/model:${PYTHONPATH:-}"32export CUBLAS_WORKSPACE_CONFIG=:4096:233export NCCL_DEBUG="${NCCL_DEBUG:-INFO}"34export PYTHONFAULTHANDLER=135 36if [[ "${OFFLINE}" == "1" ]]; then37    export HF_HUB_OFFLINE=138    export TRANSFORMERS_OFFLINE=139fi40 41if [[ -n "${CONDA_ENV}" ]] && command -v conda >/dev/null 2>&1; then42    # shellcheck disable=SC109143    source "$(conda info --base)/etc/profile.d/conda.sh"44    conda activate "${CONDA_ENV}"45fi46 47if [[ ! -f "${TOKENIZER_PATH}" ]]; then48    echo "Missing TOKENIZER_PATH: ${TOKENIZER_PATH}" >&249    exit 150fi51 52if [[ ! -f "${TRAIN_DATA_PATH}" ]]; then53    echo "Missing TRAIN_DATA_PATH: ${TRAIN_DATA_PATH}" >&254    echo "Create it first with CodonTransformer.CodonData.prepare_training_data." >&255    exit 156fi57 58args=(59    "${PROJECT_DIR}/scripts/pretrain.py"60    --tokenizer_path "${TOKENIZER_PATH}"61    --train_data_path "${TRAIN_DATA_PATH}"62    --checkpoint_dir "${CHECKPOINT_DIR}"63    --batch_size "${BATCH_SIZE}"64    --max_epochs "${MAX_EPOCHS}"65    --num_workers "${NUM_WORKERS}"66    --accumulate_grad_batches "${ACCUMULATE_GRAD_BATCHES}"67    --num_gpus "${NUM_GPUS}"68    --learning_rate "${LEARNING_RATE}"69    --warmup_fraction "${WARMUP_FRACTION}"70    --save_interval "${SAVE_INTERVAL}"71    --seed "${SEED}"72)73 74if [[ "${DEBUG}" == "1" ]]; then75    args+=(--debug)76fi77 78python "${args[@]}"79