File size: 3,510 Bytes
f2ad668 | 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 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 | import os
from functools import cache
from pathlib import Path
from huggingface_hub import snapshot_download
SEQUENCE_BOS_TOKEN = 0
SEQUENCE_PAD_TOKEN = 1
SEQUENCE_EOS_TOKEN = 2
SEQUENCE_CHAINBREAK_TOKEN = 31
SEQUENCE_MASK_TOKEN = 32
VQVAE_CODEBOOK_SIZE = 4096
VQVAE_SPECIAL_TOKENS = {
"MASK": VQVAE_CODEBOOK_SIZE,
"EOS": VQVAE_CODEBOOK_SIZE + 1,
"BOS": VQVAE_CODEBOOK_SIZE + 2,
"PAD": VQVAE_CODEBOOK_SIZE + 3,
"CHAINBREAK": VQVAE_CODEBOOK_SIZE + 4,
}
VQVAE_DIRECTION_LOSS_BINS = 16
VQVAE_PAE_BINS = 64
VQVAE_MAX_PAE_BIN = 31.0
VQVAE_PLDDT_BINS = 50
STRUCTURE_MASK_TOKEN = VQVAE_SPECIAL_TOKENS["MASK"]
STRUCTURE_BOS_TOKEN = VQVAE_SPECIAL_TOKENS["BOS"]
STRUCTURE_EOS_TOKEN = VQVAE_SPECIAL_TOKENS["EOS"]
STRUCTURE_PAD_TOKEN = VQVAE_SPECIAL_TOKENS["PAD"]
STRUCTURE_CHAINBREAK_TOKEN = VQVAE_SPECIAL_TOKENS["CHAINBREAK"]
STRUCTURE_UNDEFINED_TOKEN = 955
SASA_PAD_TOKEN = 0
SS8_PAD_TOKEN = 0
INTERPRO_PAD_TOKEN = 0
RESIDUE_PAD_TOKEN = 0
CHAIN_BREAK_STR = "|"
SEQUENCE_BOS_STR = "<cls>"
SEQUENCE_EOS_STR = "<eos>"
MASK_STR_SHORT = "_"
SEQUENCE_MASK_STR = "<mask>"
SASA_MASK_STR = "<unk>"
SS8_MASK_STR = "<unk>"
# fmt: off
SEQUENCE_VOCAB = [
"<cls>", "<pad>", "<eos>", "<unk>",
"L", "A", "G", "V", "S", "E", "R", "T", "I", "D", "P", "K",
"Q", "N", "F", "Y", "M", "H", "W", "C", "X", "B", "U", "Z",
"O", ".", "-", "|",
"<mask>",
]
# fmt: on
SEQUENCE_STANDARD_AA_MIN_TOKEN = 4 # L
SEQUENCE_STANDARD_AA_MAX_TOKEN = 24 # X (exclusive)
SSE_8CLASS_VOCAB = "GHITEBSC"
SSE_3CLASS_VOCAB = "HEC"
SSE_8CLASS_TO_3CLASS_MAP = {
"G": "H",
"H": "H",
"I": "H",
"T": "C",
"E": "E",
"B": "E",
"S": "C",
"C": "C",
}
SASA_DISCRETIZATION_BOUNDARIES = [
0.8,
4.0,
9.6,
16.4,
24.5,
32.9,
42.0,
51.5,
61.2,
70.9,
81.6,
93.3,
107.2,
125.4,
151.4,
]
MAX_RESIDUE_ANNOTATIONS = 16
TFIDF_VECTOR_SIZE = 58641
FUNCTION_TOKENS_DEPTH = 8
@staticmethod
@cache
def data_root(model: str):
if "INFRA_PROVIDER" in os.environ:
return Path("")
# Try to download from huggingface if it doesn't exist
if model.startswith("esm3"):
path = Path(snapshot_download(repo_id="biohub/esm3-sm-open-v1"))
elif model.startswith("esmc-300"):
path = Path(snapshot_download(repo_id="biohub/esmc-300m-2024-12"))
elif model.startswith("esmc-600"):
path = Path(snapshot_download(repo_id="biohub/esmc-600m-2024-12"))
elif model.startswith("esmc-6b"):
path = Path(snapshot_download(repo_id="biohub/esmc-6b-2024-12"))
else:
raise ValueError(f"{model=} is an invalid model name.")
return path
IN_REPO_DATA_FOLDER = Path(__file__).parents[2] / "data"
INTERPRO_ENTRY = IN_REPO_DATA_FOLDER / "entry_list_safety_29026.list"
INTERPRO_HIERARCHY = IN_REPO_DATA_FOLDER / "ParentChildTreeFile.txt"
INTERPRO2GO = IN_REPO_DATA_FOLDER / "ParentChildTreeFile.txt"
INTERPRO_2ID = "data/tag_dict_4_safety_filtered.json"
LSH_TABLE_PATHS = {"8bit": "data/hyperplanes_8bit_58641.npz"}
KEYWORDS_VOCABULARY = (
IN_REPO_DATA_FOLDER / "keyword_vocabulary_safety_filtered_58641.txt"
)
KEYWORDS_IDF = IN_REPO_DATA_FOLDER / "keyword_idf_safety_filtered_58641.npy"
RESID_CSV = "data/uniref90_and_mgnify90_residue_annotations_gt_1k_proteins.csv"
INTERPRO2KEYWORDS = IN_REPO_DATA_FOLDER / "interpro_29026_to_keywords_58641.csv"
|