| import torch |
| from torch.utils.data import Dataset |
| from utils import SmilesEnumerator |
| import numpy as np |
| import re |
|
|
| 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[idx], 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 |
|
|
