PerturbReason/PerturbReason_dataset_code

1#!/bin/bash2#SBATCH --job-name=nlm_05013#SBATCH --chdir=/scratch/group/PROJECT_NAME/OMICS_DATA/baseline_collections/natureLM_official4#SBATCH --output=/scratch/group/PROJECT_NAME/OMICS_DATA/baseline_collections/natureLM_official/logs/natureLM_0501_%A_%a.log5#SBATCH --partition=gpu6#SBATCH --nodes=17#SBATCH --ntasks-per-node=18#SBATCH --gres=gpu:a100:29#SBATCH --cpus-per-task=1610#SBATCH --mem=128G11#SBATCH --time=6:00:0012#SBATCH --account=ACCOUNT_ID13#SBATCH --array=0-914 15set -euo pipefail16 17SCRIPT_ROOT="/scratch/group/PROJECT_NAME/OMICS_DATA/baseline_collections/natureLM_official"18LOGS_DIR="/scratch/group/PROJECT_NAME/OMICS_DATA/baseline_collections/natureLM_official/logs"19 20echo "start time: $(date)"21 22export CACHE_ROOT="/scratch/group/PROJECT_NAME/cache"23export HF_HOME="$CACHE_ROOT/hf"24export HF_DATASETS_CACHE="$CACHE_ROOT/hf/datasets"25export XDG_CACHE_HOME="$CACHE_ROOT/xdg_cache"26export VLLM_CONFIG_ROOT="$CACHE_ROOT/vllm_config"27export TRITON_CACHE_DIR="$CACHE_ROOT/triton"28export FLASHINFER_CACHE_DIR="$CACHE_ROOT/flashinfer"29export TORCHINDUCTOR_ROOT="$CACHE_ROOT/torchinductor"30 31TASK_ID="${SLURM_ARRAY_TASK_ID:-0}"32JOB_TAG="${SLURM_JOB_ID:-manual}_${TASK_ID}"33export TMPDIR="$CACHE_ROOT/tmp/naturelm_${JOB_TAG}"34export TMP="$TMPDIR"35export TEMP="$TMPDIR"36export TORCHINDUCTOR_CACHE_DIR="$TORCHINDUCTOR_ROOT/natureLM_${JOB_TAG}"37 38mkdir -p $HF_HOME $HF_DATASETS_CACHE $XDG_CACHE_HOME $VLLM_CONFIG_ROOT $TRITON_CACHE_DIR $FLASHINFER_CACHE_DIR $TMPDIR $TORCHINDUCTOR_CACHE_DIR39 40export VLLM_HOST_IP=$(hostname -I | awk '{print $1}')41export VLLM_USE_TORCH_COMPILE=042export TORCH_COMPILE=043export TORCHINDUCTOR_VEC_ISA_OK=044export VLLM_WORKER_MULTIPROC_METHOD=spawn45export CUDA_DEVICE_ORDER=PCI_BUS_ID46 47export PATH="/scratch/user/USER_NAME/State/state/.venv/bin:$PATH"48export PYTHONPATH="$SCRIPT_ROOT${PYTHONPATH:+:$PYTHONPATH}"49export NATURELM_VLLM_MOE_FALLBACK=150PYTHON_BIN="/scratch/user/USER_NAME/State/state/.venv/bin/python"51 52MODEL_PATH="/scratch/group/PROJECT_NAME/OMICS_DATA/natureLM"53DATASET_ROOT="/scratch/group/PROJECT_NAME/OMICS_DATA/baseline_collections/dataset_0426_final"54OUTPUT_ROOT="/scratch/group/PROJECT_NAME/OMICS_DATA/baseline_collections/natureLM_official"55 56CONTEXTS=(hidden_context noisy_context)57SPLITS=(id_test cell_ood pert_ood double_ood other_test)58 59if (( TASK_ID < 0 || TASK_ID >= 10 )); then60    echo "ERROR: TASK_ID=${TASK_ID} is out of range [0, 9]"61    exit 162fi63 64CONTEXT_IDX=$(( TASK_ID / 5 ))65SPLIT_IDX=$(( TASK_ID % 5 ))66 67CONTEXT=${CONTEXTS[$CONTEXT_IDX]}68SPLIT=${SPLITS[$SPLIT_IDX]}69 70declare -A CTX_TO_OUT71CTX_TO_OUT[hidden_context]=hidden_output_050172CTX_TO_OUT[noisy_context]=noisy_output_050173 74OUTPUT_NAME=${CTX_TO_OUT[$CONTEXT]}75 76INPUT_DIR="${DATASET_ROOT}/${CONTEXT}/${SPLIT}"77OUTPUT_DIR="${OUTPUT_ROOT}/${OUTPUT_NAME}/${SPLIT}"78 79LIMIT_FILES="${LIMIT_FILES:-0}"80LIMIT_RECORDS="${LIMIT_RECORDS:-0}"81OVERWRITE_EXISTING="${OVERWRITE_EXISTING:-1}"82EXTRA_ARGS=()83if (( LIMIT_FILES > 0 )); then84    EXTRA_ARGS+=(--limit_files "$LIMIT_FILES")85fi86if (( LIMIT_RECORDS > 0 )); then87    EXTRA_ARGS+=(--limit_records "$LIMIT_RECORDS")88fi89if [[ "$OVERWRITE_EXISTING" == "1" ]]; then90    EXTRA_ARGS+=(--overwrite_existing)91fi92 93CMD=(94    "$PYTHON_BIN" "$SCRIPT_ROOT/natureLM_batch_vllm.py"95    --model_path "$MODEL_PATH"96    --input_dir "$INPUT_DIR"97    --output_dir "$OUTPUT_DIR"98    --max_new_tokens 599    --tp_size 2100    --gpu_memory_utilization 0.9101    --temperature 0.1102    --top_p 1.0103    --guided_decoding104    "${EXTRA_ARGS[@]}"105)106 107echo "=========================================="108echo "NatureLM 0501 Pipeline — guided decoding + label order shuffle"109echo "Context : $CONTEXT"110echo "Split   : $SPLIT  (array index $TASK_ID)"111echo "Model   : $MODEL_PATH"112echo "Input   : $INPUT_DIR"113echo "Output  : $OUTPUT_DIR"114echo "TMPDIR  : $TMPDIR"115if (( LIMIT_FILES > 0 || LIMIT_RECORDS > 0 )); then116    echo "Smoke   : limit_files=$LIMIT_FILES limit_records=$LIMIT_RECORDS"117fi118echo "Overwrite existing outputs: $OVERWRITE_EXISTING"119echo "CUDA_VISIBLE_DEVICES : ${CUDA_VISIBLE_DEVICES:-<unset>}"120echo "=========================================="121 122if [[ "${DRY_RUN:-0}" == "1" ]]; then123    printf 'DRY_RUN command:'124    printf ' %q' "${CMD[@]}"125    printf '\n'126    echo "Skipping execution because DRY_RUN=1"127    exit 0128fi129 130module purge131module load CUDA/12.8.0132 133nvidia-smi --query-gpu=index,uuid,name --format=csv,noheader || true134 135mkdir -p "$OUTPUT_DIR"136mkdir -p "$LOGS_DIR"137 138"${CMD[@]}"139 140echo "Finished: $CONTEXT / $SPLIT"141echo "end time: $(date)"142