File size: 1,908 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 | #!/usr/bin/env bash
set -euo pipefail
# Export a Lightning .ckpt checkpoint to a plain .pt model file for inference.
# Run this after scripts/slurm/run_finetune.sh.
# This script does not request SLURM resources.
PROJECT_DIR="${PROJECT_DIR:-/public/home/scnb9biwet/jiangqq/CodonTransformer-main}"
CONDA_ENV="${CONDA_ENV:-struct-evo}"
CHECKPOINT_DIR="${CHECKPOINT_DIR:-${PROJECT_DIR}/weight/checkpoints/finetune}"
CHECKPOINT_PATH="${CHECKPOINT_PATH:-}"
OUTPUT_MODEL_PATH="${OUTPUT_MODEL_PATH:-${CHECKPOINT_DIR}/finetuned_model.pt}"
NUM_ORGANISMS="${NUM_ORGANISMS:-164}"
cd "${PROJECT_DIR}"
export PYTHONPATH="${PROJECT_DIR}/model:${PYTHONPATH:-}"
export CHECKPOINT_DIR
export CHECKPOINT_PATH
export OUTPUT_MODEL_PATH
export NUM_ORGANISMS
export PYTHONFAULTHANDLER=1
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
if [[ -z "${CHECKPOINT_PATH}" ]]; then
CHECKPOINT_PATH="$(find "${CHECKPOINT_DIR}" -maxdepth 1 -type f -name 'epoch=*-step=*.ckpt' | sort | tail -n 1)"
fi
if [[ -z "${CHECKPOINT_PATH}" ]] || [[ ! -f "${CHECKPOINT_PATH}" ]]; then
echo "Missing CHECKPOINT_PATH. Set it explicitly or check CHECKPOINT_DIR: ${CHECKPOINT_DIR}" >&2
exit 1
fi
mkdir -p "$(dirname "${OUTPUT_MODEL_PATH}")"
export CHECKPOINT_PATH
python - <<'PY'
import os
from CodonTransformer.CodonPrediction import create_model_from_checkpoint
checkpoint_path = os.environ["CHECKPOINT_PATH"]
output_model_path = os.environ["OUTPUT_MODEL_PATH"]
num_organisms = int(os.environ["NUM_ORGANISMS"])
print(f"Checkpoint: {checkpoint_path}")
print(f"Output model: {output_model_path}")
print(f"Num organisms: {num_organisms}")
create_model_from_checkpoint(
checkpoint_dir=checkpoint_path,
output_model_dir=output_model_path,
num_organisms=num_organisms,
)
PY
|