File size: 1,976 Bytes
53e66de | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 | #!/usr/bin/env bash
set -euo pipefail
# Single deterministic inference for one protein sequence.
# Run this inside an allocated/interactive GPU session. No SLURM resources are requested here.
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}"
PROTEIN="${PROTEIN:-MFWY}"
ORGANISM="${ORGANISM:-Escherichia coli general}"
OFFLINE="${OFFLINE:-1}"
cd "${PROJECT_DIR}"
export HF_HOME
export PYTHONPATH="${PROJECT_DIR}/model:${PYTHONPATH:-}"
export PROTEIN
export ORGANISM
export OFFLINE
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
python - <<'PY'
import os
import torch
from transformers import AutoTokenizer, BigBirdForMaskedLM
from CodonTransformer.CodonJupyter import format_model_output
from CodonTransformer.CodonPrediction import predict_dna_sequence
protein = os.environ["PROTEIN"]
organism = os.environ["ORGANISM"]
local_files_only = os.environ.get("OFFLINE", "1") == "1"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"HF_HOME: {os.environ.get('HF_HOME')}")
print(f"Device: {device}")
print(f"Local files only: {local_files_only}")
tokenizer = AutoTokenizer.from_pretrained(
"adibvafa/CodonTransformer",
local_files_only=local_files_only,
)
model = BigBirdForMaskedLM.from_pretrained(
"adibvafa/CodonTransformer",
local_files_only=local_files_only,
).to(device)
output = predict_dna_sequence(
protein=protein,
organism=organism,
device=device,
tokenizer=tokenizer,
model=model,
attention_type="original_full",
deterministic=True,
)
print(format_model_output(output))
PY
|