Upload 17 files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+bin/foldseek filter=lfs diff=lfs merge=lfs -text

bin/foldseek

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0f903f1e906d2a7b38335caf2cb65323d40c0740825e2b2ab122cc7787d7e22b
+size 100165624

environment.yml

@@ -0,0 +1,191 @@
+name: pltnum
+channels:
+  - anaconda
+  - pytorch
+  - nvidia
+  - conda-forge
+  - defaults
+dependencies:
+  - _libgcc_mutex=0.1=conda_forge
+  - _openmp_mutex=4.5=2_kmp_llvm
+  - abseil-cpp=20211102.0=hd4dd3e8_0
+  - aiohttp=3.9.5=py311h5eee18b_0
+  - aiosignal=1.2.0=pyhd3eb1b0_0
+  - arrow-cpp=14.0.2=h374c478_1
+  - attrs=23.1.0=py311h06a4308_0
+  - aws-c-auth=0.6.19=h5eee18b_0
+  - aws-c-cal=0.5.20=hdbd6064_0
+  - aws-c-common=0.8.5=h5eee18b_0
+  - aws-c-compression=0.2.16=h5eee18b_0
+  - aws-c-event-stream=0.2.15=h6a678d5_0
+  - aws-c-http=0.6.25=h5eee18b_0
+  - aws-c-io=0.13.10=h5eee18b_0
+  - aws-c-mqtt=0.7.13=h5eee18b_0
+  - aws-c-s3=0.1.51=hdbd6064_0
+  - aws-c-sdkutils=0.1.6=h5eee18b_0
+  - aws-checksums=0.1.13=h5eee18b_0
+  - aws-crt-cpp=0.18.16=h6a678d5_0
+  - aws-sdk-cpp=1.10.55=h721c034_0
+  - biopython=1.84=py311h331c9d8_0
+  - blas=1.0=mkl
+  - boost-cpp=1.82.0=hdb19cb5_2
+  - bottleneck=1.3.7=py311hf4808d0_0
+  - brotli-python=1.0.9=py311h6a678d5_8
+  - bzip2=1.0.8=h5eee18b_6
+  - c-ares=1.19.1=h5eee18b_0
+  - ca-certificates=2024.7.4=hbcca054_0
+  - certifi=2024.7.4=pyhd8ed1ab_0
+  - cffi=1.16.0=py311h5eee18b_1
+  - charset-normalizer=2.0.4=pyhd3eb1b0_0
+  - cloudpickle=3.0.0=pyhd8ed1ab_0
+  - cryptography=42.0.5=py311hdda0065_1
+  - cuda-cudart=12.1.105=0
+  - cuda-cupti=12.1.105=0
+  - cuda-libraries=12.1.0=0
+  - cuda-nvrtc=12.1.105=0
+  - cuda-nvtx=12.1.105=0
+  - cuda-opencl=12.5.39=0
+  - cuda-runtime=12.1.0=0
+  - cuda-version=12.5=3
+  - datasets=2.19.1=py311h06a4308_0
+  - dill=0.3.8=py311h06a4308_0
+  - ffmpeg=4.3=hf484d3e_0
+  - filelock=3.13.1=py311h06a4308_0
+  - freetype=2.12.1=h4a9f257_0
+  - frozenlist=1.4.0=py311h5eee18b_0
+  - fsspec=2024.3.1=py311h06a4308_0
+  - gflags=2.2.2=h6a678d5_1
+  - glog=0.5.0=h6a678d5_1
+  - gmp=6.2.1=h295c915_3
+  - gmpy2=2.1.2=py311hc9b5ff0_0
+  - gnutls=3.6.15=he1e5248_0
+  - grpc-cpp=1.48.2=he1ff14a_1
+  - huggingface_hub=0.23.1=py311h06a4308_0
+  - icu=73.1=h6a678d5_0
+  - idna=3.7=py311h06a4308_0
+  - intel-openmp=2023.1.0=hdb19cb5_46306
+  - jinja2=3.1.4=py311h06a4308_0
+  - joblib=1.4.2=py311h06a4308_0
+  - jpeg=9e=h5eee18b_1
+  - krb5=1.20.1=h143b758_1
+  - lame=3.100=h7b6447c_0
+  - lcms2=2.12=h3be6417_0
+  - ld_impl_linux-64=2.38=h1181459_1
+  - lerc=3.0=h295c915_0
+  - libboost=1.82.0=h109eef0_2
+  - libbrotlicommon=1.0.9=h5eee18b_8
+  - libbrotlidec=1.0.9=h5eee18b_8
+  - libbrotlienc=1.0.9=h5eee18b_8
+  - libcublas=12.1.0.26=0
+  - libcufft=11.0.2.4=0
+  - libcufile=1.10.1.7=0
+  - libcurand=10.3.6.82=0
+  - libcurl=8.7.1=h251f7ec_0
+  - libcusolver=11.4.4.55=0
+  - libcusparse=12.0.2.55=0
+  - libdeflate=1.17=h5eee18b_1
+  - libedit=3.1.20230828=h5eee18b_0
+  - libev=4.33=h7f8727e_1
+  - libevent=2.1.12=hdbd6064_1
+  - libffi=3.4.4=h6a678d5_1
+  - libgcc-ng=14.1.0=h77fa898_0
+  - libgfortran-ng=11.2.0=h00389a5_1
+  - libgfortran5=11.2.0=h1234567_1
+  - libiconv=1.16=h5eee18b_3
+  - libidn2=2.3.4=h5eee18b_0
+  - libjpeg-turbo=2.0.0=h9bf148f_0
+  - libllvm14=14.0.6=hdb19cb5_3
+  - libnghttp2=1.57.0=h2d74bed_0
+  - libnpp=12.0.2.50=0
+  - libnvjitlink=12.1.105=0
+  - libnvjpeg=12.1.1.14=0
+  - libpng=1.6.39=h5eee18b_0
+  - libprotobuf=3.20.3=he621ea3_0
+  - libssh2=1.11.0=h251f7ec_0
+  - libstdcxx-ng=14.1.0=hc0a3c3a_0
+  - libtasn1=4.19.0=h5eee18b_0
+  - libthrift=0.15.0=h1795dd8_2
+  - libtiff=4.5.1=h6a678d5_0
+  - libunistring=0.9.10=h27cfd23_0
+  - libuuid=1.41.5=h5eee18b_0
+  - libwebp-base=1.3.2=h5eee18b_0
+  - llvm-openmp=14.0.6=h9e868ea_0
+  - llvmlite=0.43.0=py311h6a678d5_0
+  - lz4-c=1.9.4=h6a678d5_1
+  - markupsafe=2.1.3=py311h5eee18b_0
+  - mkl=2023.1.0=h213fc3f_46344
+  - mkl-service=2.4.0=py311h5eee18b_1
+  - mkl_fft=1.3.8=py311h5eee18b_0
+  - mkl_random=1.2.4=py311hdb19cb5_0
+  - mpc=1.1.0=h10f8cd9_1
+  - mpfr=4.0.2=hb69a4c5_1
+  - mpmath=1.3.0=py311h06a4308_0
+  - multidict=6.0.4=py311h5eee18b_0
+  - multiprocess=0.70.15=py311h06a4308_0
+  - ncurses=6.4=h6a678d5_0
+  - nettle=3.7.3=hbbd107a_1
+  - networkx=3.3=py311h06a4308_0
+  - numba=0.60.0=py311h4bc866e_0
+  - numexpr=2.8.7=py311h65dcdc2_0
+  - numpy=1.26.0=py311h08b1b3b_0
+  - numpy-base=1.26.0=py311hf175353_0
+  - openh264=2.1.1=h4ff587b_0
+  - openjpeg=2.4.0=h9ca470c_1
+  - openssl=3.3.1=h4bc722e_2
+  - orc=1.7.4=hb3bc3d3_1
+  - packaging=24.1=py311h06a4308_0
+  - pandas=2.1.1=py311ha02d727_0
+  - pillow=10.3.0=py311h5eee18b_0
+  - pip=24.0=py311h06a4308_0
+  - pyarrow=14.0.2=py311hb6e97c4_0
+  - pybind11-abi=4=hd3eb1b0_1
+  - pycparser=2.21=pyhd3eb1b0_0
+  - pyopenssl=24.0.0=py311h06a4308_0
+  - pysocks=1.7.1=py311h06a4308_0
+  - python=3.11.8=h955ad1f_0
+  - python-dateutil=2.9.0post0=py311h06a4308_2
+  - python-tzdata=2023.3=pyhd3eb1b0_0
+  - python-xxhash=2.0.2=py311h5eee18b_1
+  - python_abi=3.11=2_cp311
+  - pytorch=2.3.1=py3.11_cuda12.1_cudnn8.9.2_0
+  - pytorch-cuda=12.1=ha16c6d3_5
+  - pytorch-mutex=1.0=cuda
+  - pytz=2024.1=py311h06a4308_0
+  - pyyaml=6.0.1=py311h5eee18b_0
+  - re2=2022.04.01=h295c915_0
+  - readline=8.2=h5eee18b_0
+  - regex=2023.10.3=py311h5eee18b_0
+  - requests=2.31.0=py311h06a4308_0
+  - s2n=1.3.27=hdbd6064_0
+  - safetensors=0.4.2=py311h24d97f6_1
+  - scikit-learn=1.2.2=py311h6a678d5_1
+  - scipy=1.13.1=py311h08b1b3b_0
+  - setuptools=69.5.1=py311h06a4308_0
+  - shap=0.45.1=cpu_py311h9c1f9ec_0
+  - six=1.16.0=pyhd3eb1b0_1
+  - slicer=0.0.8=pyhd8ed1ab_0
+  - snappy=1.1.10=h6a678d5_1
+  - sqlite=3.45.3=h5eee18b_0
+  - sympy=1.12=py311h06a4308_0
+  - tbb=2021.8.0=hdb19cb5_0
+  - threadpoolctl=3.5.0=py311h92b7b1e_0
+  - tk=8.6.14=h39e8969_0
+  - tokenizers=0.15.1=py311h22610ee_0
+  - torchaudio=2.3.1=py311_cu121
+  - torchtriton=2.3.1=py311
+  - torchvision=0.18.1=py311_cu121
+  - tqdm=4.66.4=py311h92b7b1e_0
+  - transformers=4.38.2=pyhd8ed1ab_0
+  - typing-extensions=4.11.0=py311h06a4308_0
+  - typing_extensions=4.11.0=py311h06a4308_0
+  - tzdata=2024a=h04d1e81_0
+  - urllib3=1.26.19=py311h06a4308_0
+  - utf8proc=2.6.1=h5eee18b_1
+  - wheel=0.43.0=py311h06a4308_0
+  - xlrd=2.0.1=pyhd3eb1b0_1
+  - xxhash=0.8.0=h7f8727e_3
+  - xz=5.4.6=h5eee18b_1
+  - yaml=0.2.5=h7b6447c_0
+  - yarl=1.9.3=py311h5eee18b_0
+  - zlib=1.2.13=h5eee18b_1
+  - zstd=1.5.5=hc292b87_2

scripts/apply_foldseek_to_pdb.py

@@ -0,0 +1,61 @@
+import os
+import random
+import argparse
+import glob
+import pandas as pd
+import multiprocessing as mp
+from foldseek_util import get_struc_seq
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--pdb_dir",
+        type=str,
+        default="./pdb_files",
+        help="Directory containing PDB files.",
+    )
+    parser.add_argument(
+        "--num_processes",
+        type=int,
+        default=2,
+        help="Number of processes to use for multiprocessing. Default is 2.",
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default="./data",
+        help="Output directory.",
+    )
+    return parser.parse_args()
+
+
+def get_foldseek_seq(pdb_path):
+    parsed_seqs = get_struc_seq(
+        "bin/foldseek",
+        pdb_path,
+        ["A"],
+        process_id=random.randint(0, 10000000),
+    )["A"]
+    return parsed_seqs
+
+
+if __name__ == "__main__":
+    config = parse_args()
+
+    pdb_files = glob.glob(os.path.join(config.pdb_dir, "*.pdb"))
+
+    with mp.Pool(config.num_processes) as pool:
+        output = pool.map(get_foldseek_seq, pdb_files)
+
+    aa, foldseek, aa_foldseek = zip(*output)
+
+    result = {}
+    result["file"] = pdb_files
+    result["aa"] = aa
+    result["foldseek"] = foldseek
+    result["aa_foldseek"] = aa_foldseek
+
+    df = pd.DataFrame(result)
+
+    df.to_csv(os.path.join(config.output_dir, "foldseek_result.csv"), index=False)

scripts/augmentation.py

@@ -0,0 +1,61 @@
+import random
+
+def random_change_augmentation(aas, cfg):
+    residue_tokens = ("A", "C", "D", "E", "F", "G", "H", "I", "K", "L", "M", "N", "P", "Q", "R", "S", "T", "V", "W", "Y")
+    stracture_aware_tokens = ("a", "c", "d", "e", "f", "g", "h", "i", "k", "l", "m", "n", "p", "q", "r", "s", "t", "v", "w", "y")
+    length = len(aas)
+    swap_indices = random.sample(
+        range(length), int(length * cfg.random_change_ratio)
+    )
+    new_aas = ""
+    for i, aa in enumerate(aas):
+        if i in swap_indices:
+            if aas[i] in residue_tokens:
+                new_aas += random.choice(residue_tokens)
+            elif aas[i] in stracture_aware_tokens:
+                new_aas += random.choice(stracture_aware_tokens)
+        else:
+            new_aas += aa
+    return new_aas
+
+
+def mask_augmentation(aas, cfg):
+    length = len(aas)
+    swap_indices = random.sample(
+        range(0, length // cfg.token_length),
+        int(length // cfg.token_length * cfg.mask_ratio),
+    )
+    for ith in swap_indices:
+        aas = (
+            aas[: ith * cfg.token_length]
+            + "@" * cfg.token_length
+            + aas[(ith + 1) * cfg.token_length :]
+        )
+    aas = aas.replace("@@", "<mask>").replace("@", "<mask>")
+    return aas
+
+
+def random_delete_augmentation(aas, cfg):
+    length = len(aas)
+    swap_indices = random.sample(
+        range(0, length // cfg.token_length),
+        int(length // cfg.token_length * cfg.random_delete_ratio),
+    )
+    for ith in swap_indices:
+        aas = (
+            aas[: ith * cfg.token_length]
+            + "@" * cfg.token_length
+            + aas[(ith + 1) * cfg.token_length :]
+        )
+    aas = aas.replace("@@", "").replace("@", "")
+    return aas
+
+
+def truncate_augmentation(aas, cfg):
+    length = len(aas)
+    if length > cfg.max_length:
+        diff = length - cfg.max_length
+        start = random.randint(0, diff)
+        return aas[start : start + cfg.max_length]
+    else:
+        return aas

scripts/calculate_shap.py

@@ -0,0 +1,160 @@
+import os
+import glob
+import sys
+import argparse
+import pandas as pd
+import torch
+from transformers import AutoTokenizer
+import shap
+
+sys.path.append(".")
+from utils import seed_everything, save_pickle
+from models import PLTNUM, PLTNUM_PreTrainedModel
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Calculate SHAP values with a pretrained protein half-life prediction model."
+    )
+    parser.add_argument(
+        "--data_path",
+        type=str,
+        required=True,
+        help="Path to the input data.",
+    )
+    parser.add_argument(
+        "--model",
+        type=str,
+        default="westlake-repl/SaProt_650M_AF2",
+        help="Pretrained model name or path.",
+    )
+    parser.add_argument(
+        "--architecture",
+        type=str,
+        default="SaProt",
+        help="Model architecture: 'ESM2', 'SaProt', or 'LSTM'.",
+    )
+    parser.add_argument(
+        "--folds",
+        type=int,
+        default=10,
+        help="The number of folds for prediction.",
+    )
+    parser.add_argument(
+        "--do_cross_validation",
+        action="store_true",
+        default=False,
+        help="Use cross validation for prediction. If True, you have to specify the 'data_path' that contanins fold information, 'folds' for the number of folds, and 'model_path' for the directory of the model weights.",
+    )
+    parser.add_argument(
+        "--model_path",
+        type=str,
+        required=False,
+        help="Path to the model weight(s).",
+    )
+    parser.add_argument("--batch_size", type=int, default=4, help="Batch size.")
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=42,
+        help="Seed for reproducibility.",
+    )
+    parser.add_argument(
+        "--max_length",
+        type=int,
+        default=512,
+        help="Maximum input sequence length. Two tokens are used fo <cls> and <eos> tokens. So the actual length of input sequence is max_length - 2. Padding or truncation is applied to make the length of input sequence equal to max_length.",
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default="./output",
+        help="Output directory.",
+    )
+    parser.add_argument(
+        "--task",
+        type=str,
+        default="classification",
+        help="Task type: 'classification' or 'regression'.",
+    )
+    parser.add_argument(
+        "--sequence_col",
+        type=str,
+        default="aa_foldseek",
+        help="Column name fot the input sequence.",
+    )
+    parser.add_argument(
+        "--max_evals",
+        type=int,
+        default=5000,
+        help="Number of evaluations for SHAP values calculation.",
+    )
+
+
+    return parser.parse_args()
+
+
+def calculate_shap_fn(texts, model, cfg):
+    if len(texts) == 1:
+        texts = texts[0]
+    else:
+        texts = texts.tolist()
+
+    inputs = cfg.tokenizer(
+        texts,
+        return_tensors="pt",
+        padding=True,
+        truncation=True,
+        max_length=cfg.max_length,
+    )
+    inputs = {k: v.to(cfg.device) for k, v in inputs.items()}
+    with torch.no_grad():
+        outputs = model(inputs)
+        outputs = torch.sigmoid(outputs).detach().cpu().numpy()
+    return outputs
+
+
+if __name__ == "__main__":
+    config = parse_args()
+    config.device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    if not os.path.exists(config.output_dir):
+        os.makedirs(config.output_dir)
+    seed_everything(config.seed)
+
+    df = pd.read_csv(config.data_path)
+    config.tokenizer = AutoTokenizer.from_pretrained(config.model)
+
+    if config.do_cross_validation:
+        model_weights = glob.glob(os.path.join(config.model_path, "*.pth"))
+        for fold in range(config.folds):
+            model = PLTNUM(config).to(config.device)
+            model_weight = [w for w in model_weights if f"fold{fold}.pth" in w][0]
+            model.load_state_dict(torch.load(model_weight, map_location="cpu"))
+            model.eval()
+
+            df_fold = df[df["fold"] == fold].reset_index(drop=True)
+            explainer = shap.Explainer(lambda x: calculate_shap_fn(x, model, config), config.tokenizer)
+            shap_values = explainer(
+                df_fold[config.sequence_col].values.tolist(),
+                batch_size=config.batch_size,
+                max_evals=config.max_evals,
+            )
+
+            save_pickle(os.path.join(config.output_dir, f"shap_values_fold{fold}.pickle"), shap_values)
+    else:
+        model = PLTNUM_PreTrainedModel.from_pretrained(config.model_path, cfg=config).to(config.device)
+        model.eval()
+
+        # build an explainer using a token masker
+        explainer = shap.Explainer(lambda x: calculate_shap_fn(x, model, config), config.tokenizer)
+
+        shap_values = explainer(
+            df[config.sequence_col].values.tolist(),
+            batch_size=config.batch_size,
+            max_evals=config.max_evals,
+        )
+
+        save_pickle(
+            os.path.join(config.output_dir, "shap_values.pickle"), shap_values
+        )

scripts/convert_to_PreTrainedModel.py

@@ -0,0 +1,61 @@
+import os
+import sys
+import argparse
+
+import torch
+import shutil
+
+# Append the utils module path
+sys.path.append("../")
+from models import PLTNUM
+
+
+def parse_args():
+    """
+    Parse command line arguments.
+    """
+    parser = argparse.ArgumentParser(
+        description="Convert the model implemented with nn.Module to a model implemented with transformers' PreTrainedModel."
+    )
+    parser.add_argument(
+        "--model_path",
+        type=str,
+        help="The path to a model weight which you want to convert.",
+    )
+    parser.add_argument(
+        "--config_and_tokenizer_path",
+        type=str,
+        help="The path to a config and tokenizer of the model which you want to convert.",
+    )
+    parser.add_argument(
+        "--model",
+        type=str,
+        help="The name of the base model of the finetuned model",
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default="./",
+        help="Directory to save the prediction.",
+    )
+    parser.add_argument(
+        "--task",
+        type=str,
+        default="classification",
+    )
+
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    config = parse_args()
+
+    if not os.path.exists(config.output_dir):
+        os.makedirs(config.output_dir)
+
+    model = PLTNUM(config)
+    model.load_state_dict(torch.load(config.model_path, map_location="cpu"))
+
+    torch.save(model.state_dict(), os.path.join(config.output_dir, "pytorch_model.bin"))
+    for filename in ["config.json", "special_tokens_map.json", "tokenizer_config.json", "vocab.txt"]:
+        shutil.copy(os.path.join(config.config_and_tokenizer_path, filename), os.path.join(config.output_dir, filename))

scripts/datasets.py

@@ -0,0 +1,127 @@
+import random
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+from augmentation import (
+    mask_augmentation,
+    random_change_augmentation,
+    random_delete_augmentation,
+    truncate_augmentation,
+)
+
+
+def tokenize_input(cfg, text):
+    inputs = cfg.tokenizer(
+        text,
+        add_special_tokens=True,
+        max_length=cfg.max_length,
+        padding="max_length",
+        truncation=True,
+        return_offsets_mapping=False,
+        return_attention_mask=True,
+    )
+    for k, v in inputs.items():
+        inputs[k] = torch.tensor(v, dtype=torch.long)
+    return inputs
+
+
+def one_hot_encoding(aa, amino_acids, cfg):
+    aa = aa[: cfg.max_length].ljust(cfg.max_length, " ")
+    one_hot = np.zeros((len(aa), len(amino_acids)))
+    for i, a in enumerate(aa):
+        if a in amino_acids:
+            one_hot[i, amino_acids.index(a)] = 1
+    return one_hot
+
+
+def one_hot_encode_input(text, cfg):
+    inputs = one_hot_encoding(text, ("A","C","D","E","F","G","H","I","K","L","M","N","P","Q","R","S","T","V","W","Y"," "), cfg)
+    return torch.tensor(inputs, dtype=torch.float)
+
+
+class PLTNUMDataset(Dataset):
+    def __init__(self, cfg, df, train=True):
+        self.df = df
+        self.cfg = cfg
+        self.train = train
+
+    def __len__(self):
+        return len(self.df)
+
+    def __getitem__(self, idx):
+        data = self.df.iloc[idx]
+        aas = self._adjust_sequence_length(data[self.cfg.sequence_col])
+
+        if self.train:
+            aas = self._apply_augmentation(aas)
+
+        aas = aas.replace("__", "<pad>")
+
+        inputs = tokenize_input(self.cfg, aas)
+
+        if "target" in data:
+            return inputs, torch.tensor(data["target"], dtype=torch.float32)
+        return inputs, np.nan
+
+    def _adjust_sequence_length(self, aas):
+        max_length = (self.cfg.max_length - 2) * self.cfg.token_length
+        if len(aas) > max_length:
+            if self.cfg.used_sequence == "left":
+                return aas[: max_length]
+            elif self.cfg.used_sequence == "right":
+                return aas[-max_length:]
+            elif self.cfg.used_sequence == "both":
+                half_max_len = max_length // 2
+                return aas[:half_max_len] + "__" + aas[-half_max_len:]
+            elif self.cfg.used_sequence == "internal":
+                offset = (len(aas) - max_length) // 2
+                return aas[offset:offset + max_length]
+        return aas
+
+    def _apply_augmentation(self, aas):
+        if self.cfg.random_change_ratio > 0:
+            aas = random_change_augmentation(aas, self.cfg)
+        if (
+            random.random() <= self.cfg.random_delete_prob
+        ) and self.cfg.random_delete_ratio > 0:
+            aas = random_delete_augmentation(aas, self.cfg)
+        if (random.random() <= self.cfg.mask_prob) and self.cfg.mask_ratio > 0:
+            aas = mask_augmentation(aas, self.cfg)
+        if random.random() <= self.cfg.truncate_augmentation_prob:
+            aas = truncate_augmentation(aas, self.cfg)
+        return aas
+
+
+class LSTMDataset(Dataset):
+    def __init__(self, cfg, df, train=True):
+        self.df = df
+        self.cfg = cfg
+        self.train = train
+
+    def __len__(self):
+        return len(self.df)
+
+    def __getitem__(self, idx):
+        data = self.df.iloc[idx]
+        aas = data[self.cfg.sequence_col]
+        aas = self._adjust_sequence_length(aas)
+        aas = aas.replace("__", "<pad>")
+
+        inputs = one_hot_encode_input(aas, self.cfg)
+
+        return inputs, torch.tensor(data["target"], dtype=torch.float32)
+
+    def _adjust_sequence_length(self, aas):
+        max_length = (self.cfg.max_length - 2) * self.cfg.token_length
+        if len(aas) > max_length:
+            if self.cfg.used_sequence == "left":
+                return aas[:max_length]
+            elif self.cfg.used_sequence == "right":
+                return aas[-max_length:]
+            elif self.cfg.used_sequence == "both":
+                half_max_len = max_length // 2
+                return aas[:half_max_len] + "__" + aas[-half_max_len:]
+            elif self.cfg.used_sequence == "internal":
+                offset = (len(aas) - max_length) // 2
+                return aas[offset:offset + max_length]
+        return aas

scripts/foldseek_util.py

@@ -0,0 +1,106 @@
+# https://github.com/westlake-repl/SaProt/blob/main/utils/foldseek_util.py
+
+# MIT License
+
+# Copyright (c) 2023 westlake-repl
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+import os
+import json
+import numpy as np
+import sys
+
+sys.path.append(".")
+
+
+# Get structural seqs from pdb file
+def get_struc_seq(
+    foldseek,
+    path,
+    chains: list = None,
+    process_id: int = 0,
+    plddt_path: str = None,
+    plddt_threshold: float = 70.0,
+) -> dict:
+    """
+
+    Args:
+        foldseek: Binary executable file of foldseek
+        path: Path to pdb file
+        chains: Chains to be extracted from pdb file. If None, all chains will be extracted.
+        process_id: Process ID for temporary files. This is used for parallel processing.
+        plddt_path: Path to plddt file. If None, plddt will not be used.
+        plddt_threshold: Threshold for plddt. If plddt is lower than this value, the structure will be masked.
+
+    Returns:
+        seq_dict: A dict of structural seqs. The keys are chain IDs. The values are tuples of
+        (seq, struc_seq, combined_seq).
+    """
+    assert os.path.exists(foldseek), f"Foldseek not found: {foldseek}"
+    assert os.path.exists(path), f"Pdb file not found: {path}"
+    assert plddt_path is None or os.path.exists(
+        plddt_path
+    ), f"Plddt file not found: {plddt_path}"
+
+    tmp_save_path = f"get_struc_seq_{process_id}.tsv"
+    cmd = f"{foldseek} structureto3didescriptor -v 0 --threads 1 --chain-name-mode 1 {path} {tmp_save_path}"
+    os.system(cmd)
+
+    seq_dict = {}
+    name = os.path.basename(path)
+    with open(tmp_save_path, "r") as r:
+        for i, line in enumerate(r):
+            desc, seq, struc_seq = line.split("\t")[:3]
+
+            # Mask low plddt
+            if plddt_path is not None:
+                with open(plddt_path, "r") as r:
+                    plddts = np.array(json.load(r)["confidenceScore"])
+
+                    # Mask regions with plddt < threshold
+                    indices = np.where(plddts < plddt_threshold)[0]
+                    np_seq = np.array(list(struc_seq))
+                    np_seq[indices] = "#"
+                    struc_seq = "".join(np_seq)
+
+            name_chain = desc.split(" ")[0]
+            chain = name_chain.replace(name, "").split("_")[-1]
+
+            if chains is None or chain in chains:
+                if chain not in seq_dict:
+                    combined_seq = "".join(
+                        [a + b.lower() for a, b in zip(seq, struc_seq)]
+                    )
+                    seq_dict[chain] = (seq, struc_seq, combined_seq)
+
+    os.remove(tmp_save_path)
+    os.remove(tmp_save_path + ".dbtype")
+    return seq_dict
+
+
+if __name__ == "__main__":
+    foldseek = "/sujin/bin/foldseek"
+    # test_path = "/sujin/Datasets/PDB/all/6xtd.cif"
+    test_path = "/sujin/Datasets/FLIP/meltome/af2_structures/A0A061ACX4.pdb"
+    plddt_path = "/sujin/Datasets/FLIP/meltome/af2_plddts/A0A061ACX4.json"
+    res = get_struc_seq(
+        foldseek, test_path, plddt_path=plddt_path, plddt_threshold=70.0
+    )
+    print(res["A"][1].lower())

scripts/get_aa_from_uniprot_accession.py

@@ -0,0 +1,93 @@
+import multiprocessing as mp
+import requests as r
+import argparse
+from Bio import SeqIO
+from io import StringIO
+import pandas as pd
+import numpy as np
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--file_path",
+        type=str,
+        required=True,
+        help="Path to the file that have a column cotaining uniprotid information.",
+    )
+    parser.add_argument(
+        "--sheet_name",
+        type=str,
+        default="Sheet1",
+        help="Name of the sheet to read. Default is Sheet1.",
+    )
+    parser.add_argument(
+        "--uniprotid_column",
+        type=str,
+        help="Name of the column that have uniprotid information. Default is None.",
+    )
+    parser.add_argument(
+        "--uniprotids_column",
+        type=str,
+        help="Name of the column that have uniprotids information. Default is None. The ids are expected to be separated by semi-colon, and the first id is used.",
+    )
+    parser.add_argument(
+        "--num_processes",
+        type=int,
+        default=2,
+        help="Number of processes to use.",
+    )
+    return parser.parse_args()
+
+
+def fetch_sequence(row, cfg):
+    try:
+        baseURL = "http://www.uniprot.org/uniprot/"
+        uniprot_id = row[cfg.uniprotid_column]
+        URL = baseURL + uniprot_id + ".fasta"
+        response = r.post(URL)
+        Data = "".join(response.text)
+        Seq = StringIO(Data)
+        pSeq = list(SeqIO.parse(Seq, "fasta"))
+        return str(pSeq[0].seq)
+    except:
+        return None
+
+
+def process_rows(df_chunk, cfg):
+    return [fetch_sequence(row, cfg) for idx, row in df_chunk.iterrows()]
+
+
+if __name__ == "__main__":
+    config = parse_args()
+
+
+    if config.file_path.endswith(".xls"):
+        df = pd.read_excel(
+            config.file_path,
+            sheet_name=config.sheet_name,
+        )
+    else:
+        df = pd.read_csv(config.file_path)
+
+    if config.uniprotid_column is None and config.uniprotids_column is None:
+        raise ValueError(
+            "Either uniprotid_column or uniprotids_column should be provided."
+        )
+    if config.uniprotids_column is not None:
+        df = df.dropna(subset=[config.uniprotids_column]).reset_index(drop=True)
+        # use the first id and ignore the subunit and domain information
+        df["uniprotid"] = df[config.uniprotids_column].apply(
+            lambda x: x.split(";")[0].split("-")[0]
+        )
+        config.uniprotid_column = "uniprotid"
+
+    df_split = np.array_split(df, config.num_processes)
+
+    with mp.Pool(processes=config.num_processes) as pool:
+        results = pool.map(lambda x: process_rows(x, config), df_split)
+
+    aas = [seq for result in results for seq in result]
+
+    df["aa"] = aas
+    df.to_csv(f"{config.file_path.split('.')[0]}_with_aa.csv", index=False)

scripts/models.py

@@ -0,0 +1,120 @@
+import torch.nn as nn
+from transformers import AutoModel, AutoConfig, PreTrainedModel 
+
+
+class PLTNUM(nn.Module):
+    def __init__(self, cfg):
+        super(PLTNUM, self).__init__()
+        self.cfg = cfg
+        self.config = AutoConfig.from_pretrained(cfg.model, output_hidden_states=True)
+        # self.model = AutoModel.from_pretrained(cfg.model, config=self.config)
+        self.model = AutoModel.from_config(config=self.config)
+
+        self.fc_dropout1 = nn.Dropout(0.8)
+        self.fc_dropout2 = nn.Dropout(0.4 if cfg.task == "classification" else 0.8)
+        self.fc = nn.Linear(self.config.hidden_size, 1)
+        self._init_weights(self.fc)
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
+            if module.bias is not None:
+                nn.init.constant_(module.bias, 0)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
+            if module.padding_idx is not None:
+                nn.init.constant_(module.weight[module.padding_idx], 0.0)
+        elif isinstance(module, nn.LayerNorm):
+            nn.init.constant_(module.bias, 0)
+            nn.init.constant_(module.weight, 1.0)
+
+    def forward(self, inputs):
+        outputs = self.model(**inputs)
+        last_hidden_state = outputs.last_hidden_state[:, 0]
+        output = (
+            self.fc(self.fc_dropout1(last_hidden_state))
+            + self.fc(self.fc_dropout2(last_hidden_state))
+        ) / 2
+        return output
+
+    def create_embedding(self, inputs):
+        outputs = self.model(**inputs)
+        last_hidden_state = outputs.last_hidden_state[:, 0]
+        return last_hidden_state
+
+
+class PLTNUM_PreTrainedModel(PreTrainedModel):
+    config_class = AutoConfig
+
+    def __init__(self, config, cfg):
+        super(PLTNUM_PreTrainedModel, self).__init__(config)
+        self.cfg = cfg
+        self.model = AutoModel.from_pretrained(self.config._name_or_path)
+
+        self.fc_dropout1 = nn.Dropout(0.8)
+        self.fc_dropout2 = nn.Dropout(0.4 if cfg.task == "classification" else 0.8)
+        self.fc = nn.Linear(self.config.hidden_size, 1)
+        self._init_weights(self.fc)
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
+            if module.bias is not None:
+                nn.init.constant_(module.bias, 0)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
+            if module.padding_idx is not None:
+                nn.init.constant_(module.weight[module.padding_idx], 0.0)
+        elif isinstance(module, nn.LayerNorm):
+            nn.init.constant_(module.bias, 0)
+            nn.init.constant_(module.weight, 1.0)
+
+    def forward(self, inputs):
+        outputs = self.model(**inputs)
+        last_hidden_state = outputs.last_hidden_state[:, 0]
+        output = (
+            self.fc(self.fc_dropout1(last_hidden_state))
+            + self.fc(self.fc_dropout2(last_hidden_state))
+        ) / 2
+        return output
+
+    def create_embedding(self, inputs):
+        outputs = self.model(**inputs)
+        last_hidden_state = outputs.last_hidden_state[:, 0]
+        return last_hidden_state
+
+
+class LSTMModel(nn.Module):
+    def __init__(self, cfg):
+        super(LSTMModel, self).__init__()
+        self.cfg = cfg
+        self.lstm = nn.LSTM(
+            input_size=21,
+            hidden_size=256,
+            num_layers=2,
+            batch_first=True,
+            bidirectional=True,
+            dropout=0.2,
+        )
+        self.fc_dropout = nn.Dropout(0.8)
+        self.fc = nn.Linear(256 * 2, 1)
+        self._init_weights(self.fc)
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
+            if module.bias is not None:
+                nn.init.constant_(module.bias, 0)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
+            if module.padding_idx is not None:
+                nn.init.constant_(module.weight[module.padding_idx], 0.0)
+        elif isinstance(module, nn.LayerNorm):
+            nn.init.constant_(module.bias, 0)
+            nn.init.constant_(module.weight, 1.0)
+
+    def forward(self, inputs):
+        outputs, _ = self.lstm(inputs)
+        last_hidden_state = outputs[:, -1, :]
+        output = self.fc(self.fc_dropout(last_hidden_state))
+        return output

scripts/predict.py

@@ -0,0 +1,168 @@
+import gc
+import os
+import sys
+import argparse
+
+import pandas as pd
+import torch
+from torch.utils.data import DataLoader
+from transformers import AutoTokenizer
+
+sys.path.append(".")
+from utils import seed_everything
+from models import PLTNUM
+from datasets import PLTNUMDataset
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Prediction script for protein sequence classification/regression."
+    )
+    parser.add_argument(
+        "--data_path",
+        type=str,
+        required=True,
+        help="Path to the input data.",
+    )
+    parser.add_argument(
+        "--model",
+        type=str,
+        default="westlake-repl/SaProt_650M_AF2",
+        help="Pretrained model name or path.",
+    )
+    parser.add_argument(
+        "--architecture",
+        type=str,
+        default="SaProt",
+        help="Model architecture: 'ESM2', 'SaProt', or 'LSTM'.",
+    )
+    parser.add_argument(
+        "--model_path",
+        type=str,
+        required=True,
+        help="Path to the model for prediction.",
+    )
+    parser.add_argument("--batch_size", type=int, default=4, help="Batch size.")
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=42,
+        help="Seed for reproducibility.",
+    )
+    parser.add_argument(
+        "--use_amp",
+        action="store_true",
+        default=False,
+        help="Use AMP for mixed precision prediction.",
+    )
+    parser.add_argument(
+        "--num_workers",
+        type=int,
+        default=4,
+        help="Number of workers for data loading.",
+    )
+    parser.add_argument(
+        "--max_length",
+        type=int,
+        default=512,
+        help="Maximum input sequence length. Two tokens are used fo <cls> and <eos> tokens. So the actual length of input sequence is max_length - 2. Padding or truncation is applied to make the length of input sequence equal to max_length.",
+    )
+    parser.add_argument(
+        "--used_sequence",
+        type=str,
+        default="left",
+        help="Which part of the sequence to use: 'left', 'right', 'both', or 'internal'.",
+    )
+    parser.add_argument(
+        "--padding_side",
+        type=str,
+        default="right",
+        help="Padding side: 'right' or 'left'.",
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default="./output",
+        help="Output directory.",
+    )
+    parser.add_argument(
+        "--task",
+        type=str,
+        default="classification",
+        help="Task type: 'classification' or 'regression'.",
+    )
+    parser.add_argument(
+        "--sequence_col",
+        type=str,
+        default="aa_foldseek",
+        help="Column name fot the input sequence.",
+    )
+
+    return parser.parse_args()
+
+
+def predict_fn(valid_loader, model, cfg):
+    model.eval()
+    predictions = []
+
+    for inputs, _ in valid_loader:
+        inputs = inputs.to(cfg.device)
+        with torch.no_grad():
+            with torch.cuda.amp.autocast(enabled=cfg.use_amp):
+                preds = (
+                    torch.sigmoid(model(inputs))
+                    if cfg.task == "classification"
+                    else model(inputs)
+                )
+        predictions += preds.cpu().tolist()
+
+    return predictions
+
+
+def predict(folds, model_path, cfg):
+    dataset = PLTNUMDataset(cfg, folds, train=False)
+    loader = DataLoader(
+        dataset,
+        batch_size=cfg.batch_size,
+        shuffle=False,
+        num_workers=cfg.num_workers,
+        pin_memory=True,
+        drop_last=False,
+    )
+
+    model = PLTNUM(cfg)
+    model.load_state_dict(torch.load(model_path, map_location=cfg.device))
+    model.to(cfg.device)
+
+    predictions = predict_fn(loader, model, cfg)
+
+    folds["raw prediction values"] = predictions
+    if cfg.task == "classification":
+        folds["binary prediction values"] = [1 if x > 0.5 else 0 for x in predictions]
+    torch.cuda.empty_cache()
+    gc.collect()
+    return folds
+
+
+if __name__ == "__main__":
+    config = parse_args()
+    config.token_length = 2 if config.architecture == "SaProt" else 1
+    config.device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    if not os.path.exists(config.output_dir):
+        os.makedirs(config.output_dir)
+
+    if config.used_sequence == "both":
+        config.max_length += 1
+
+    seed_everything(config.seed)
+
+    df = pd.read_csv(config.data_path)
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        config.model, padding_side=config.padding_side
+    )
+    config.tokenizer = tokenizer
+
+    result = predict(df, config.model_path, config)
+    result.to_csv(os.path.join(config.output_dir, "result.csv"), index=False)

scripts/predict_with_PreTrainedModel.py

@@ -0,0 +1,149 @@
+import gc
+import os
+import sys
+import argparse
+
+import pandas as pd
+import torch
+from torch.utils.data import DataLoader
+from transformers import AutoTokenizer
+
+sys.path.append(".")
+from utils import seed_everything
+from models import PLTNUM_PreTrainedModel
+from datasets import PLTNUMDataset
+from predict import predict_fn
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Prediction script for protein sequence classification/regression."
+    )
+    parser.add_argument(
+        "--data_path",
+        type=str,
+        required=True,
+        help="Path to the input data.",
+    )
+    parser.add_argument(
+        "--model",
+        type=str,
+        default="westlake-repl/SaProt_650M_AF2",
+        help="Pretrained model name or path.",
+    )
+    parser.add_argument(
+        "--architecture",
+        type=str,
+        default="SaProt",
+        help="Model architecture: 'ESM2', 'SaProt', or 'LSTM'.",
+    )
+    parser.add_argument(
+        "--model_path",
+        type=str,
+        required=True,
+        help="Path to the model for prediction.",
+    )
+    parser.add_argument("--batch_size", type=int, default=4, help="Batch size.")
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=42,
+        help="Seed for reproducibility.",
+    )
+    parser.add_argument(
+        "--use_amp",
+        action="store_true",
+        default=False,
+        help="Use AMP for mixed precision prediction.",
+    )
+    parser.add_argument(
+        "--num_workers",
+        type=int,
+        default=4,
+        help="Number of workers for data loading.",
+    )
+    parser.add_argument(
+        "--max_length",
+        type=int,
+        default=512,
+        help="Maximum input sequence length. Two tokens are used fo <cls> and <eos> tokens. So the actual length of input sequence is max_length - 2. Padding or truncation is applied to make the length of input sequence equal to max_length.",
+    )
+    parser.add_argument(
+        "--used_sequence",
+        type=str,
+        default="left",
+        help="Which part of the sequence to use: 'left', 'right', 'both', or 'internal'.",
+    )
+    parser.add_argument(
+        "--padding_side",
+        type=str,
+        default="right",
+        help="Padding side: 'right' or 'left'.",
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default="./output",
+        help="Output directory.",
+    )
+    parser.add_argument(
+        "--task",
+        type=str,
+        default="classification",
+        help="Task type: 'classification' or 'regression'.",
+    )
+    parser.add_argument(
+        "--sequence_col",
+        type=str,
+        default="aa_foldseek",
+        help="Column name fot the input sequence.",
+    )
+
+    return parser.parse_args()
+
+
+def predict(folds, model_path, cfg):
+    dataset = PLTNUMDataset(cfg, folds, train=False)
+    loader = DataLoader(
+        dataset,
+        batch_size=cfg.batch_size,
+        shuffle=False,
+        num_workers=cfg.num_workers,
+        pin_memory=True,
+        drop_last=False,
+    )
+
+    model = PLTNUM_PreTrainedModel.from_pretrained(model_path, cfg=cfg)
+    # model.load_state_dict(torch.load(os.path.join(model_path, "pytorch_model.bin"), map_location=cfg.device))
+    model.to(cfg.device)
+
+    predictions = predict_fn(loader, model, cfg)
+
+    folds["prediction"] = predictions
+    torch.cuda.empty_cache()
+    gc.collect()
+    return folds
+
+
+if __name__ == "__main__":
+    config = parse_args()
+    config.token_length = 2 if config.architecture == "SaProt" else 1
+    config.device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    if not os.path.exists(config.output_dir):
+        os.makedirs(config.output_dir)
+
+    if config.used_sequence == "both":
+        config.max_length += 1
+
+    seed_everything(config.seed)
+
+    df = pd.read_csv(config.data_path)
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        config.model_path, padding_side=config.padding_side
+    )
+    config.tokenizer = tokenizer
+
+    result = predict(df, config.model_path, config)
+    result.to_csv(os.path.join(config.output_dir, "result.csv"), index=False)

scripts/train.py

@@ -0,0 +1,493 @@
+import gc
+import os
+import sys
+import time
+import argparse
+
+import numpy as np
+import pandas as pd
+import torch
+import torch.nn as nn
+from sklearn.metrics import accuracy_score, f1_score, r2_score
+from sklearn.model_selection import StratifiedKFold
+from torch.optim import Adam
+from torch.optim.lr_scheduler import CosineAnnealingLR
+from torch.utils.data import DataLoader
+from transformers import AutoTokenizer, get_cosine_schedule_with_warmup
+
+sys.path.append(".")
+from utils import AverageMeter, get_logger, seed_everything, timeSince
+from datasets import PLTNUMDataset, LSTMDataset
+from models import PLTNUM, LSTMModel
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+print("device:", device)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Training script for protein half-life prediction."
+    )
+    parser.add_argument(
+        "--data_path",
+        type=str,
+        required=True,
+        help="Path to the training data.",
+    )
+    parser.add_argument(
+        "--model",
+        type=str,
+        default="westlake-repl/SaProt_650M_AF2",
+        help="Pretrained model name or path.",
+    )
+    parser.add_argument(
+        "--architecture",
+        type=str,
+        default="SaProt",
+        help="Model architecture: 'ESM2', 'SaProt', or 'LSTM'.",
+    )
+    parser.add_argument("--lr", type=float, default=2e-5, help="Learning rate.")
+    parser.add_argument(
+        "--epochs",
+        type=int,
+        default=5,
+        help="Number of training epochs.",
+    )
+    parser.add_argument("--batch_size", type=int, default=4, help="Batch size.")
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=42,
+        help="Seed for reproducibility.",
+    )
+    parser.add_argument(
+        "--use_amp",
+        action="store_true",
+        default=False,
+        help="Use AMP for mixed precision training.",
+    )
+    parser.add_argument(
+        "--num_workers",
+        type=int,
+        default=4,
+        help="Number of workers for data loading.",
+    )
+    parser.add_argument(
+        "--max_length",
+        type=int,
+        default=512,
+        help="Maximum input sequence length. Two tokens are used fo <cls> and <eos> tokens. So the actual length of input sequence is max_length - 2. Padding or truncation is applied to make the length of input sequence equal to max_length.",
+    )
+    parser.add_argument(
+        "--used_sequence",
+        type=str,
+        default="left",
+        help="Which part of the sequence to use: 'left', 'right', 'both', or 'internal'.",
+    )
+    parser.add_argument(
+        "--padding_side",
+        type=str,
+        default="right",
+        help="Padding side: 'right' or 'left'.",
+    )
+    parser.add_argument(
+        "--mask_ratio",
+        type=float,
+        default=0.05,
+        help="Ratio of mask tokens for augmentation.",
+    )
+    parser.add_argument(
+        "--mask_prob",
+        type=float,
+        default=0.2,
+        help="Probability to apply mask augmentation",
+    )
+    parser.add_argument(
+        "--random_delete_ratio",
+        type=float,
+        default=0.1,
+        help="Ratio of deleting tokens in augmentation.",
+    )
+    parser.add_argument(
+        "--random_delete_prob",
+        type=float,
+        default=-1,
+        help="Probability to apply random delete augmentation.",
+    )
+    parser.add_argument(
+        "--random_change_ratio",
+        type=float,
+        default=0,
+        help="Ratio of changing tokens in augmentation.",
+    )
+    parser.add_argument(
+        "--truncate_augmentation_prob",
+        type=float,
+        default=-1,
+        help="Probability to apply truncate augmentation.",
+    )
+    parser.add_argument(
+        "--n_folds",
+        type=int,
+        default=10,
+        help="Number of folds for cross-validation.",
+    )
+    parser.add_argument(
+        "--print_freq",
+        type=int,
+        default=300,
+        help="Log print frequency.",
+    )
+    parser.add_argument(
+        "--freeze_layer",
+        type=int,
+        default=-1,
+        help="Freeze layers of the model. -1 means no layers are frozen.",
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default="./output",
+        help="Output directory.",
+    )
+    parser.add_argument(
+        "--task",
+        type=str,
+        default="classification",
+        help="Task type: 'classification' or 'regression'.",
+    )
+    parser.add_argument(
+        "--target_col",
+        type=str,
+        default="Protein half-life average [h]",
+        help="Column name of the target.",
+    )
+    parser.add_argument(
+        "--sequence_col",
+        type=str,
+        default="aa_foldseek",
+        help="Column name fot the input sequence.",
+    )
+
+    return parser.parse_args()
+
+
+def train_fn(train_loader, model, criterion, optimizer, epoch, cfg):
+    model.train()
+    scaler = torch.cuda.amp.GradScaler(enabled=cfg.use_amp)
+    losses = AverageMeter()
+    label_list, pred_list = [], []
+    start = time.time()
+
+    for step, (inputs, labels) in enumerate(train_loader):
+        inputs, labels = inputs.to(cfg.device), labels.to(cfg.device)
+        labels = (
+            labels.float()
+            if cfg.task == "classification"
+            else labels.to(dtype=torch.half)
+        )
+        batch_size = labels.size(0)
+
+        with torch.cuda.amp.autocast(enabled=cfg.use_amp):
+            y_preds = model(inputs)
+        loss = criterion(y_preds, labels.view(-1, 1))
+        losses.update(loss.item(), batch_size)
+
+        scaler.scale(loss).backward()
+        scaler.step(optimizer)
+        scaler.update()
+        optimizer.zero_grad()
+
+        label_list += labels.tolist()
+        pred_list += y_preds.tolist()
+
+        if step % cfg.print_freq == 0 or step == len(train_loader) - 1:
+            if cfg.task == "classification":
+                pred_list_new = (torch.Tensor(pred_list) > 0.5).to(dtype=torch.long)
+                acc = accuracy_score(label_list, pred_list_new > 0.5)
+                cfg.logger.info(
+                    f"Epoch: [{epoch + 1}][{step}/{len(train_loader)}] "
+                    f"Elapsed {timeSince(start, float(step + 1) / len(train_loader))} "
+                    f"Loss: {losses.val:.4f}({losses.avg:.4f}) "
+                    f"LR: {optimizer.param_groups[0]['lr']:.8f} "
+                    f"Accuracy: {acc:.4f}"
+                )
+            elif cfg.task == "regression":
+                r2 = r2_score(label_list, pred_list)
+                cfg.logger.info(
+                    f"Epoch: [{epoch + 1}][{step}/{len(train_loader)}] "
+                    f"Elapsed {timeSince(start, float(step + 1) / len(train_loader))} "
+                    f"Loss: {losses.val:.4f}({losses.avg:.4f}) "
+                    f"R2 Score: {r2:.4f} "
+                    f"LR: {optimizer.param_groups[0]['lr']:.8f}"
+                )
+    if cfg.task == "classification":
+        pred_list_new = (torch.Tensor(pred_list) > 0.5).to(dtype=torch.long)
+        acc = accuracy_score(label_list, pred_list_new)
+        return losses.avg, acc
+    elif cfg.task == "regression":
+        return losses.avg, r2_score(label_list, pred_list)
+
+
+def valid_fn(valid_loader, model, criterion, cfg):
+    losses = AverageMeter()
+    model.eval()
+    label_list, pred_list = [], []
+    start = time.time()
+
+    for step, (inputs, labels) in enumerate(valid_loader):
+        inputs, labels = inputs.to(cfg.device), labels.to(cfg.device)
+        labels = (
+            labels.float()
+            if cfg.task == "classification"
+            else labels.to(dtype=torch.half)
+        )
+
+        with torch.no_grad():
+            with torch.cuda.amp.autocast(enabled=cfg.use_amp):
+                y_preds = (
+                    torch.sigmoid(model(inputs))
+                    if cfg.task == "classification"
+                    else model(inputs)
+                )
+        loss = criterion(y_preds, labels.view(-1, 1))
+        losses.update(loss.item(), labels.size(0))
+
+        label_list += labels.tolist()
+        pred_list += y_preds.tolist()
+
+        if step % cfg.print_freq == 0 or step == len(valid_loader) - 1:
+            if cfg.task == "classification":
+                pred_list_new = (torch.Tensor(pred_list) > 0.5).to(dtype=torch.long)
+                acc = accuracy_score(label_list, pred_list_new > 0.5)
+                f1 = f1_score(label_list, pred_list_new, average="macro")
+                cfg.logger.info(
+                    f"EVAL: [{step}/{len(valid_loader)}] "
+                    f"Elapsed {timeSince(start, float(step + 1) / len(valid_loader))} "
+                    f"Loss: {losses.val:.4f}({losses.avg:.4f}) "
+                    f"Accuracy: {acc:.4f} "
+                    f"F1 Score: {f1:.4f}"
+                )
+            elif cfg.task == "regression":
+                r2 = r2_score(label_list, pred_list)
+                cfg.logger.info(
+                    f"EVAL: [{step}/{len(valid_loader)}] "
+                    f"Elapsed {timeSince(start, float(step + 1) / len(valid_loader))} "
+                    f"Loss: {losses.val:.4f}({losses.avg:.4f}) "
+                    f"R2 Score: {r2:.4f}"
+                )
+
+    if cfg.task == "classification":
+        pred_list_new = (torch.Tensor(pred_list) > 0.5).to(dtype=torch.long)
+        return (
+            f1_score(label_list, pred_list_new, average="macro"),
+            accuracy_score(label_list, pred_list_new),
+            pred_list,
+        )
+    elif cfg.task == "regression":
+        return losses.avg, r2_score(label_list, pred_list), np.array(pred_list)
+
+
+def train_loop(folds, fold, cfg):
+    cfg.logger.info(f"================== fold: {fold} training ======================")
+    train_folds = folds[folds["fold"] != fold].reset_index(drop=True)
+    valid_folds = folds[folds["fold"] == fold].reset_index(drop=True)
+
+    if cfg.architecture in ["ESM2", "SaProt"]:
+        train_dataset = PLTNUMDataset(cfg, train_folds, train=True)
+        valid_dataset = PLTNUMDataset(cfg, valid_folds, train=False)
+    elif cfg.architecture == "LSTM":
+        train_dataset = LSTMDataset(cfg, train_folds, train=True)
+        valid_dataset = LSTMDataset(cfg, valid_folds, train=False)
+
+    train_loader = DataLoader(
+        train_dataset,
+        batch_size=cfg.batch_size,
+        shuffle=True,
+        num_workers=cfg.num_workers,
+        pin_memory=True,
+        drop_last=True,
+    )
+    valid_loader = DataLoader(
+        valid_dataset,
+        batch_size=cfg.batch_size,
+        shuffle=False,
+        num_workers=cfg.num_workers,
+        pin_memory=True,
+        drop_last=False,
+    )
+
+    if cfg.architecture in ["ESM2", "SaProt"]:
+        model = PLTNUM(cfg)
+        if cfg.freeze_layer >= 0:
+            for name, param in model.named_parameters():
+                if f"model.encoder.layer.{cfg.freeze_layer}" in name:
+                    break
+                param.requires_grad = False
+        model.config.save_pretrained(cfg.output_dir)
+    elif cfg.architecture == "LSTM":
+        model = LSTMModel(cfg)
+
+    model.to(cfg.device)
+
+    optimizer = Adam(model.parameters(), lr=cfg.lr)
+    if cfg.architecture in ["ESM2", "SaProt"]:
+        scheduler = CosineAnnealingLR(
+            optimizer,
+            **{"T_max": 2, "eta_min": 1.0e-6, "last_epoch": -1},
+        )
+    elif cfg.architecture == "LSTM":
+        scheduler = get_cosine_schedule_with_warmup(
+            optimizer, num_warmup_steps=0, num_training_steps=cfg.epochs, num_cycles=0.5
+        )
+
+    criterion = nn.BCEWithLogitsLoss() if cfg.task == "classification" else nn.MSELoss()
+    best_score = 0 if cfg.task == "classification" else float("inf")
+
+    for epoch in range(cfg.epochs):
+        start_time = time.time()
+        # train
+        avg_loss, train_score = train_fn(
+            train_loader, model, criterion, optimizer, epoch, cfg
+        )
+        scheduler.step()
+
+        # eval
+        val_score, val_score2, predictions = valid_fn(
+            valid_loader, model, criterion, cfg
+        )
+
+        elapsed = time.time() - start_time
+
+        if cfg.task == "classification":
+            cfg.logger.info(
+                f"Epoch {epoch+1} - avg_train_loss: {avg_loss:.4f}  train_acc: {train_score:.4f}  valid_acc: {val_score2:.4f}  valid_f1: {val_score:.4f}  time: {elapsed:.0f}s"
+            )
+        elif cfg.task == "regression":
+            cfg.logger.info(
+                f"Epoch {epoch+1} - avg_train_loss: {avg_loss:.4f}  train_r2: {train_score:.4f}  valid_r2: {val_score2:.4f}  valid_loss: {val_score:.4f}  time: {elapsed:.0f}s"
+            )
+
+        if (cfg.task == "classification" and best_score < val_score) or (
+            cfg.task == "regression" and best_score > val_score
+        ):
+            best_score = val_score
+            cfg.logger.info(f"Epoch {epoch+1} - Save Best Score: {val_score:.4f} Model")
+            torch.save(
+                predictions,
+                os.path.join(cfg.output_dir, f"predictions.pth"),
+            )
+            torch.save(
+                model.state_dict(),
+                os.path.join(cfg.output_dir, f"model_fold{fold}.pth"),
+            )
+
+    predictions = torch.load(
+        os.path.join(cfg.output_dir, f"predictions.pth"), map_location="cpu"
+    )
+    valid_folds["prediction"] = predictions
+    cfg.logger.info(f"[Fold{fold}] Best score: {best_score}")
+    torch.cuda.empty_cache()
+    gc.collect()
+    return valid_folds
+
+
+def get_embedding(folds, fold, path, cfg):
+    valid_folds = folds[folds["fold"] == fold].reset_index(drop=True)
+    valid_dataset = PLTNUMDataset(cfg, valid_folds, train=False)
+
+    valid_loader = DataLoader(
+        valid_dataset,
+        batch_size=cfg.batch_size,
+        shuffle=False,
+        num_workers=cfg.num_workers,
+        pin_memory=True,
+        drop_last=False,
+    )
+
+    model = PLTNUM(cfg)
+    model.load_state_dict(torch.load(path, map_location=torch.device("cpu")))
+    model.to(device)
+
+    model.eval()
+    embedding_list = []
+    for inputs, _ in valid_loader:
+        inputs = inputs.to(device)
+        with torch.no_grad():
+            with torch.cuda.amp.autocast(enabled=cfg.use_amp):
+                embedding = model.create_embedding(inputs)
+        embedding_list += embedding.tolist()
+
+    torch.cuda.empty_cache()
+    gc.collect()
+    return embedding_list
+
+
+if __name__ == "__main__":
+    config = parse_args()
+    config.token_length = 2 if config.architecture == "SaProt" else 1
+    config.device = device
+
+    if not os.path.exists(config.output_dir):
+        os.makedirs(config.output_dir)
+
+    if config.used_sequence == "both":
+        config.max_length += 1
+
+    LOGGER = get_logger(os.path.join(config.output_dir, "output"))
+    config.logger = LOGGER
+
+    seed_everything(config.seed)
+
+    train_df = (
+        pd.read_csv(config.data_path)
+        .drop_duplicates(subset=[config.sequence_col], keep="first")
+        .reset_index(drop=True)
+    )
+    train_df["T1/2 [h]"] = train_df[config.target_col]
+
+    if config.task == "classification":
+        train_df["target"] = (
+            train_df["T1/2 [h]"] > np.median(train_df["T1/2 [h]"])
+        ).astype(int)
+        train_df["class"] = train_df["target"]
+    elif config.task == "regression":
+        train_df["log1p(T1/2 [h])"] = np.log1p(train_df["T1/2 [h]"])
+        train_df["log1p(T1/2 [h])"] = (
+            train_df["log1p(T1/2 [h])"] - min(train_df["log1p(T1/2 [h])"])
+        ) / (max(train_df["log1p(T1/2 [h])"]) - min(train_df["log1p(T1/2 [h])"]))
+        train_df["target"] = train_df["log1p(T1/2 [h])"]
+
+        def get_class(row, class_num=5):
+            denom = 1 / class_num
+            num = row["log1p(T1/2 [h])"]
+            for target in range(class_num):
+                if denom * target <= num and num < denom * (target + 1):
+                    break
+            row["class"] = target
+            return row
+
+        train_df = train_df.apply(get_class, axis=1)
+
+    train_df["fold"] = -1
+    kf = StratifiedKFold(
+        n_splits=config.n_folds, shuffle=True, random_state=config.seed
+    )
+    for fold, (trn_ind, val_ind) in enumerate(kf.split(train_df, train_df["class"])):
+        train_df.loc[val_ind, "fold"] = int(fold)
+
+    if config.architecture in ["ESM2", "SaProt"]:
+        tokenizer = AutoTokenizer.from_pretrained(
+            config.model, padding_side=config.padding_side
+        )
+        tokenizer.save_pretrained(config.output_dir)
+        config.tokenizer = tokenizer
+
+    oof_df = pd.DataFrame()
+    for fold in range(config.n_folds):
+        _oof_df = train_loop(train_df, fold, config)
+        oof_df = pd.concat([oof_df, _oof_df], axis=0)
+
+    oof_df = oof_df.reset_index(drop=True)
+    oof_df.to_csv(os.path.join(config.output_dir, "oof_df.csv"), index=False)

scripts/use_foldseek_for_uniprot.py

@@ -0,0 +1,106 @@
+import os
+import random
+import argparse
+import pandas as pd
+import multiprocessing as mp
+from foldseek_util import get_struc_seq
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--file_path",
+        type=str,
+        required=True,
+        help="Path to the file containing uniprotid information.",
+    )
+    parser.add_argument(
+        "--sheet_name",
+        type=str,
+        default="Sheet1",
+        help="Name of the sheet to read (for Excel files). Default is 'Sheet1'.",
+    )
+    parser.add_argument(
+        "--pdb_dir",
+        type=str,
+        default="pdb_files/UP000000589_10090_MOUSE_v4",
+        help="Directory containing PDB files.",
+    )
+    parser.add_argument(
+        "--uniprotid_column",
+        type=str,
+        help="Name of the column containing UniprotID information.",
+    )
+    parser.add_argument(
+        "--uniprotids_column",
+        type=str,
+        help="Name of the column containing multiple UniprotIDs (separated by semicolons). The first ID will be used.",
+    )
+    parser.add_argument(
+        "--num_processes",
+        type=int,
+        default=2,
+        help="Number of processes to use for multiprocessing. Default is 2.",
+    )
+    return parser.parse_args()
+
+
+def validate_columns(cfg, df):
+    if cfg.uniprotid_column is None and cfg.uniprotids_column is None:
+        raise ValueError("Either --uniprotid_column or --uniprotids_column must be provided.")
+    if cfg.uniprotids_column:
+        df = df.dropna(subset=[cfg.uniprotids_column]).reset_index(drop=True)
+        df["uniprotid"] = df[cfg.uniprotids_column].apply(lambda x: x.split(";")[0].split("-")[0])
+        cfg.uniprotid_column = "uniprotid"
+    return df.dropna(subset=[cfg.uniprotid_column]).reset_index(drop=True)
+
+
+def find_pdb_files(pdb_dir, uniprot_ids):
+    pdf_files = os.listdir(pdb_dir)
+    pdb_paths = []
+    for uniprot_id in uniprot_ids:
+        matches = [pdf_file for pdf_file in sorted(pdf_files) if uniprot_id in pdf_file]
+        pdb_paths.append(matches[0] if matches else None)
+    return pdb_paths
+
+
+def get_foldseek_seq(pdb_path, cfg):
+    parsed_seqs = get_struc_seq(
+        "bin/foldseek",
+        os.path.join(cfg.pdb_dir, pdb_path),
+        ["A"],
+        process_id=random.randint(0, 10000000),
+    )["A"]
+    return parsed_seqs
+
+
+if __name__ == "__main__":
+        
+    config = parse_args()
+
+    if config.file_path.endswith(".xls") or config.file_path.endswith(".xlsx"):
+        df = pd.read_excel(
+            config.file_path,
+            sheet_name=config.sheet_name,
+        )
+    else:
+        df = pd.read_csv(config.file_path)
+    df = validate_columns(config, df)
+
+    df = df.dropna(subset=[config.uniprotid_column]).reset_index(drop=True)
+
+    uniprot_ids = df[config.uniprotid_column].tolist()
+    pdb_paths = find_pdb_files(config.pdb_dir, uniprot_ids)
+    df["pdb_path"] = pdb_paths
+    df = df.dropna(subset=["pdb_path"]).reset_index(drop=True)
+    df = df.drop_duplicates(subset=[config.uniprotid_column]).reset_index(drop=True)
+
+    with mp.Pool(config.num_processes) as pool:
+        output = pool.map(lambda x: get_foldseek_seq(x, config), df["pdb_path"].tolist())
+
+    aa, foldseek, aa_foldseek = zip(*output)
+
+    df["aa"] = aa
+    df["foldseek"] = foldseek
+    df["aa_foldseek"] = aa_foldseek
+    df.to_csv(f"{config.file_path.split('.')[0]}_foldseek.csv", index=False)

utils.py

@@ -0,0 +1,88 @@
+import random
+import os
+import math
+import time
+import numpy as np
+import pickle
+import torch
+import logging
+
+
+def get_logger(filename: str):
+    """Creates and returns a logger that logs to both the console and a file."""
+    logger = logging.getLogger(__name__)
+    logger.setLevel(logging.INFO)
+
+    # Console handler
+    stream_handler = logging.StreamHandler()
+    stream_handler.setFormatter(logging.Formatter("%(message)s"))
+    logger.addHandler(stream_handler)
+
+    # File handler
+    file_handler = logging.FileHandler(f"{filename}.log")
+    file_handler.setFormatter(logging.Formatter("%(message)s"))
+    logger.addHandler(file_handler)
+
+    return logger
+
+
+def seed_everything(seed: int):
+    """Sets random seed for reproducibility across various libraries."""
+    random.seed(seed)
+    os.environ["PYTHONHASHSEED"] = str(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+
+class AverageMeter:
+    """Tracks and stores the average and current values."""
+
+    def __init__(self):
+        self.reset()
+
+    def reset(self):
+        self.val = 0
+        self.avg = 0
+        self.sum = 0
+        self.count = 0
+
+    def update(self, val, n=1):
+        self.val = val
+        self.sum += val * n
+        self.count += n
+        self.avg = self.sum / self.count
+
+
+def as_minutes(s: int) -> str:
+    """Converts seconds to a string in minutes and seconds."""
+    m = math.floor(s / 60)
+    s -= m * 60
+    return "%dm %ds" % (m, s)
+
+
+def timeSince(since: float, percent: float) -> str:
+    now = time.time()
+    s = now - since
+    es = s / (percent)
+    rs = es - s
+    return "%s (remain %s)" % (as_minutes(s), as_minutes(rs))
+
+
+def convert_all_1d(array: list) -> list:
+    """Converts 0-dimensional arrays in a list to 1-dimensional arrays."""
+    return [np.array([item]) if item.ndim == 0 else item for item in array]
+
+
+def save_pickle(path: str, contents):
+    """Saves contents to a pickle file."""
+    with open(path, "wb") as f:
+        pickle.dump(contents, f)
+
+
+def load_pickle(path: str):
+    """Loads contents from a pickle file."""
+    with open(path, "rb") as f:
+        return pickle.load(f)