OneScience-Group/CodonTransformer

1#!/usr/bin/env bash2set -euo pipefail3 4# Single deterministic inference with an exported finetuned .pt model.5# Run this after scripts/slurm/export_finetuned_model.sh.6# This script does not request SLURM resources.7 8PROJECT_DIR="${PROJECT_DIR:-/public/home/scnb9biwet/jiangqq/CodonTransformer-main}"9CONDA_ENV="${CONDA_ENV:-struct-evo}"10 11TOKENIZER_PATH="${TOKENIZER_PATH:-${PROJECT_DIR}/model/src/CodonTransformerTokenizer.json}"12MODEL_PATH="${MODEL_PATH:-${PROJECT_DIR}/weight/checkpoints/finetune/finetuned_model.pt}"13PROTEIN="${PROTEIN:-MFWY}"14ORGANISM="${ORGANISM:-Escherichia coli general}"15MATCH_PROTEIN="${MATCH_PROTEIN:-1}"16 17cd "${PROJECT_DIR}"18 19export PYTHONPATH="${PROJECT_DIR}/model:${PYTHONPATH:-}"20export TOKENIZER_PATH21export MODEL_PATH22export PROTEIN23export ORGANISM24export MATCH_PROTEIN25export PYTHONFAULTHANDLER=126 27if [[ -n "${CONDA_ENV}" ]] && command -v conda >/dev/null 2>&1; then28    # shellcheck disable=SC109129    source "$(conda info --base)/etc/profile.d/conda.sh"30    conda activate "${CONDA_ENV}"31fi32 33if [[ ! -f "${TOKENIZER_PATH}" ]]; then34    echo "Missing TOKENIZER_PATH: ${TOKENIZER_PATH}" >&235    exit 136fi37 38if [[ ! -f "${MODEL_PATH}" ]]; then39    echo "Missing MODEL_PATH: ${MODEL_PATH}" >&240    exit 141fi42 43python - <<'PY'44import os45 46import torch47 48from CodonTransformer.CodonJupyter import format_model_output49from CodonTransformer.CodonPrediction import predict_dna_sequence50 51device = torch.device("cuda" if torch.cuda.is_available() else "cpu")52match_protein = os.environ.get("MATCH_PROTEIN", "1") == "1"53 54print(f"Device: {device}")55print(f"Tokenizer: {os.environ['TOKENIZER_PATH']}")56print(f"Model: {os.environ['MODEL_PATH']}")57 58output = predict_dna_sequence(59    protein=os.environ["PROTEIN"],60    organism=os.environ["ORGANISM"],61    device=device,62    tokenizer=os.environ["TOKENIZER_PATH"],63    model=os.environ["MODEL_PATH"],64    attention_type="original_full",65    deterministic=True,66    match_protein=match_protein,67)68 69print(format_model_output(output))70PY71