File size: 825 Bytes
3b99abb | 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 | import itertools
aa_set = {"A", "C", "D", "E", "F", "G", "H", "I", "K", "L", "M", "N", "P", "Q", "R", "S", "T", "V", "W", "Y"}
aa_list = ["A", "C", "D", "E", "F", "G", "H", "I", "K", "L", "M", "N", "P", "Q", "R", "S", "T", "V", "W", "Y"]
foldseek_seq_vocab = "ACDEFGHIKLMNPQRSTVWY#"
foldseek_struc_vocab = "pynwrqhgdlvtmfsaeikc#"
struc_unit = "abcdefghijklmnopqrstuvwxyz"
def create_vocab(size: int) -> dict:
"""
Args:
size: Size of the vocabulary
Returns:
vocab: Vocabulary
"""
token_len = 1
while size > len(struc_unit) ** token_len:
token_len += 1
vocab = {}
for i, token in enumerate(itertools.product(struc_unit, repeat=token_len)):
vocab[i] = "".join(token)
if len(vocab) == size:
vocab[i+1] = "#"
return vocab |