-
Notifications
You must be signed in to change notification settings - Fork 0
/
dataset.py
77 lines (53 loc) · 2.47 KB
/
dataset.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
# This code was adapted from LigGPT https://github.com/devalab/molgpt
# with modifications.
import torch
from torch.utils.data import Dataset
from utils import SmilesEnumerator
import numpy as np
import re
import math
class SmileDataset(Dataset):
def __init__(self, args, data, content, block_size, aug_prob = 0.5, prop = None, scaffold = None, scaffold_maxlen = None):
chars = sorted(list(set(content)))
data_size, vocab_size = len(data), len(chars)
print('data has %d smiles, %d unique characters.' % (data_size, vocab_size))
self.stoi = { ch:i for i,ch in enumerate(chars) }
self.itos = { i:ch for i,ch in enumerate(chars) }
self.max_len = block_size
self.vocab_size = vocab_size
self.data = data
self.prop = prop
self.sca = scaffold
self.scaf_max_len = scaffold_maxlen
self.debug = args.debug
self.tfm = SmilesEnumerator()
self.aug_prob = aug_prob
def __len__(self):
if self.debug:
return math.ceil(len(self.data) / (self.max_len + 1))
else:
return len(self.data)
def __getitem__(self, idx):
smiles, prop, scaffold = self.data[idx], self.prop.iloc[idx,:].values, self.sca[idx]
smiles = smiles.strip()
scaffold = scaffold.strip()
p = np.random.uniform()
if p < self.aug_prob:
smiles = self.tfm.randomize_smiles(smiles)
pattern = "(\[[^\]]+]|<|Br?|Cl?|N|O|S|P|F|I|b|c|n|o|s|p|\(|\)|\.|=|#|-|\+|\\\\|\/|:|~|@|\?|>|\*|\$|\%[0-9]{2}|[0-9])"
regex = re.compile(pattern)
smiles += str('<')*(self.max_len - len(regex.findall(smiles)))
if len(regex.findall(smiles)) > self.max_len:
smiles = smiles[:self.max_len]
smiles=regex.findall(smiles)
scaffold += str('<')*(self.scaf_max_len - len(regex.findall(scaffold)))
if len(regex.findall(scaffold)) > self.scaf_max_len:
scaffold = scaffold[:self.scaf_max_len]
scaffold=regex.findall(scaffold)
dix = [self.stoi[s] for s in smiles]
sca_dix = [self.stoi[s] for s in scaffold]
sca_tensor = torch.tensor(sca_dix, dtype=torch.long)
x = torch.tensor(dix[:-1], dtype=torch.long)
y = torch.tensor(dix[1:], dtype=torch.long)
prop = torch.tensor([prop], dtype = torch.float)
return x, y, prop, sca_tensor