-
Notifications
You must be signed in to change notification settings - Fork 7
/
smiles_model.py
78 lines (65 loc) · 2.6 KB
/
smiles_model.py
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
#from utils import *
import re
import torch
from tqdm import tqdm
def smiles_atom_tokenizer (smi):
pattern = "(\[[^\]]+]|Br?|Cl?|N|O|S|P|F|I|b|c|n|o|s|p|\(|\)|\.|=|#|-|\+|\\\\|\/|:|~|@|\?|>|\*|\$|\%[0-9]{2}|[0-9])"
regex = re.compile(pattern)
tokens = [token for token in regex.findall(smi)]
return tokens
def reverse_vocab (vocab):
return dict((v,k) for k,v in vocab.items())
class SmilesModel ():
def __init__ (self, smiles, naug=3):
if type(smiles) is str:
smiles = extract_smiles(smiles)
self.special_tokens = ['<pad>', '<unk>', '<bos>', '<eos>', '>', '.']
vocab = set()
for i in tqdm(range(len(smiles))):
sm = smiles[i]
for n in range(naug):
tokens = smiles_atom_tokenizer(augment_smile(sm))
vocab |= set(tokens)
nspec = len(self.special_tokens)
self.vocab = dict(zip(sorted(vocab),
range(nspec, nspec+len(vocab))))
for i,spec in enumerate(self.special_tokens):
self.vocab[spec] = i
self.rev_vocab = reverse_vocab(self.vocab)
self.vocsize = len(self.vocab)
def encode (self, seq):
return [self.vocab[token] for token in smiles_atom_tokenizer(seq)]
def decode (self, seq):
seq = [s for s in seq if s >= 4]
return "".join([self.rev_vocab[code] for code in seq])
class RawSmilesModel ():
def __init__ (self, smiles):
self.special_tokens = ['<pad>', '<unk>', '<bos>', '<eos>', '>', '.']
vocab = set()
for i in tqdm(range(len(smiles))):
sm = smiles[i]
chars = set(list(sm))
chars -= set(self.special_tokens)
vocab |= chars
nspec = len(self.special_tokens)
self.vocab = dict(zip(sorted(vocab),
range(nspec, nspec+len(vocab))))
for i,spec in enumerate(self.special_tokens):
self.vocab[spec] = i
self.rev_vocab = reverse_vocab(self.vocab)
self.vocsize = len(self.vocab)
def encode (self, seq):
return [self.vocab[char] for char in seq]
def decode (self, seq):
seq = [s for s in seq if s >= 4]
return "".join([self.rev_vocab[code] for code in seq])
if __name__ == "__main__":
smiles = extract_smiles("../data/maximal.csv")
print('Building Raw Smiles model')
sys.stdout.flush()
RSM = RawSmilesModel(smiles)
torch.save(RSM, "raw_smiles_model.pt")
print('Building Smiles model')
sys.stdout.flush()
SM = SmilesModel(smiles)
torch.save(SM, "smiles_model.pt")