OneScience-Group/CodonTransformer

1#!/usr/bin/env bash2set -euo pipefail3 4# Multiple variable inference for one protein sequence.5# Run this inside an allocated/interactive GPU session. No SLURM resources are requested here.6 7PROJECT_DIR="${PROJECT_DIR:-/public/home/scnb9biwet/jiangqq/CodonTransformer-main}"8HF_HOME="${HF_HOME:-/public/home/scnb9biwet/.cache/huggingface}"9CONDA_ENV="${CONDA_ENV:-struct-evo}"10 11PROTEIN="${PROTEIN:-MFWY}"12ORGANISM="${ORGANISM:-Escherichia coli general}"13OUTPUT_CSV="${OUTPUT_CSV:-${PROJECT_DIR}/outputs/multiple_predictions.csv}"14NUM_SEQUENCES="${NUM_SEQUENCES:-5}"15TEMPERATURE="${TEMPERATURE:-0.5}"16TOP_P="${TOP_P:-0.95}"17MATCH_PROTEIN="${MATCH_PROTEIN:-1}"18OFFLINE="${OFFLINE:-1}"19 20cd "${PROJECT_DIR}"21mkdir -p "$(dirname "${OUTPUT_CSV}")"22 23export HF_HOME24export PYTHONPATH="${PROJECT_DIR}/model:${PYTHONPATH:-}"25export PROTEIN26export ORGANISM27export OUTPUT_CSV28export NUM_SEQUENCES29export TEMPERATURE30export TOP_P31export MATCH_PROTEIN32export OFFLINE33export PYTHONFAULTHANDLER=134 35if [[ "${OFFLINE}" == "1" ]]; then36    export HF_HUB_OFFLINE=137    export TRANSFORMERS_OFFLINE=138fi39 40if [[ -n "${CONDA_ENV}" ]] && command -v conda >/dev/null 2>&1; then41    # shellcheck disable=SC109142    source "$(conda info --base)/etc/profile.d/conda.sh"43    conda activate "${CONDA_ENV}"44fi45 46python - <<'PY'47import os48 49import pandas as pd50import torch51from transformers import AutoTokenizer, BigBirdForMaskedLM52 53from CodonTransformer.CodonJupyter import format_model_output54from CodonTransformer.CodonPrediction import predict_dna_sequence55 56protein = os.environ["PROTEIN"]57organism = os.environ["ORGANISM"]58output_csv = os.environ["OUTPUT_CSV"]59num_sequences = int(os.environ["NUM_SEQUENCES"])60temperature = float(os.environ["TEMPERATURE"])61top_p = float(os.environ["TOP_P"])62match_protein = os.environ.get("MATCH_PROTEIN", "1") == "1"63local_files_only = os.environ.get("OFFLINE", "1") == "1"64 65device = torch.device("cuda" if torch.cuda.is_available() else "cpu")66print(f"HF_HOME: {os.environ.get('HF_HOME')}")67print(f"Device: {device}")68print(f"Local files only: {local_files_only}")69print(70    "Sampling: "71    f"num_sequences={num_sequences}, temperature={temperature}, "72    f"top_p={top_p}, match_protein={match_protein}"73)74 75tokenizer = AutoTokenizer.from_pretrained(76    "adibvafa/CodonTransformer",77    local_files_only=local_files_only,78)79model = BigBirdForMaskedLM.from_pretrained(80    "adibvafa/CodonTransformer",81    local_files_only=local_files_only,82).to(device)83 84outputs = predict_dna_sequence(85    protein=protein,86    organism=organism,87    device=device,88    tokenizer=tokenizer,89    model=model,90    attention_type="original_full",91    deterministic=False,92    temperature=temperature,93    top_p=top_p,94    num_sequences=num_sequences,95    match_protein=match_protein,96)97 98rows = []99for index, output in enumerate(outputs, start=1):100    print(f"\nSequence {index}")101    print(format_model_output(output))102    rows.append(103        {104            "sequence_index": index,105            "protein_sequence": output.protein,106            "organism": output.organism,107            "predicted_dna": output.predicted_dna,108        }109    )110 111pd.DataFrame(rows).to_csv(output_csv, index=False)112print(f"Saved predictions to {output_csv}")113PY114