staplebridge/data/vocab.py

AMINO_ACIDS = list("ACDEFGHIKLMNPQRSTVWY")
NCAA_TOKENS = ["X", "B"]
PAD_TOKEN = "<pad>"
UNK_TOKEN = "<unk>"

ALL_TOKENS = [PAD_TOKEN, UNK_TOKEN] + AMINO_ACIDS + NCAA_TOKENS
TOKEN_TO_ID = {t: i for i, t in enumerate(ALL_TOKENS)}
ID_TO_TOKEN = {i: t for t, i in TOKEN_TO_ID.items()}

HYDROPHOBIC = set("AILMFWVY")
RISKY_TOKENS = {"W", "F", "B"}