OneScience-Group/ESM

1#!/usr/bin/env bash2set -euo pipefail3 4SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"5PROJECT_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)"6cd "${PROJECT_ROOT}"7 8ESM_WEIGHT_DIR="${ESM_WEIGHT_DIR:-${PROJECT_ROOT}/weight}"9ESM_OUTPUT_DIR="${ESM_OUTPUT_DIR:-${PROJECT_ROOT}/outputs}"10ESM_FASTA="${ESM_FASTA:-${PROJECT_ROOT}/data/fasta/few_proteins.fasta}"11ESM2_8M_WEIGHT="${ESM2_8M_WEIGHT:-${ESM_WEIGHT_DIR}/esm2_t6_8M_UR50D.pt}"12 13mkdir -p "${ESM_OUTPUT_DIR}"14 15if [[ ! -f "${ESM2_8M_WEIGHT}" ]]; then16  echo "Missing ${ESM2_8M_WEIGHT}"17  echo "Run: bash scripts/download_weights.sh ${ESM_WEIGHT_DIR}"18  exit 119fi20 21python scripts/extract.py \22  "${ESM2_8M_WEIGHT}" \23  "${ESM_FASTA}" \24  "${ESM_OUTPUT_DIR}/embeddings" \25  --include mean per_tok \26  --repr_layers 627 28if [[ "${RUN_ESMFOLD:-0}" == "1" ]]; then29  python scripts/fold.py \30    -i "${ESM_FASTA}" \31    -o "${ESM_OUTPUT_DIR}/pdb" \32    --model-dir "${ESM_WEIGHT_DIR}" \33    --cpu-only34fi35 36if [[ "${RUN_VARIANT_PREDICTION:-0}" == "1" ]]; then37  python scripts/variant_prediction/predict.py \38    --model-location "${ESM_WEIGHT_DIR}/esm1v_t33_650M_UR90S_1.pt" \39    --sequence "${ESM_VARIANT_SEQUENCE:?Set ESM_VARIANT_SEQUENCE}" \40    --dms-input data/variant_prediction/BLAT_ECOLX_Ranganathan2015.csv \41    --mutation-col mutant \42    --dms-output "${ESM_OUTPUT_DIR}/variant_prediction.csv" \43    --offset-idx 24 \44    --scoring-strategy wt-marginals45fi46