LLM4Mol/dataloader.py at main · ChnQ/LLM4Mol · GitHub

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
import json
import torch
from torch.utils.data import Dataset as TorchDataset


# ChatGPT outputs
class MolDataset(torch.utils.data.Dataset):
    def __init__(self, dataset_path, encodings):
        self.encodings = encodings
        with open(dataset_path, 'r', encoding='UTF-8') as f:
            self.all_infos = json.load(f)

    def __getitem__(self, idx):
        item = {key: torch.tensor(val[idx])
                for key, val in self.encodings.items()}
        if self.all_infos[idx]['label'] == '1':
            label = 1
        else:
            label = 0
        item['labels'] = torch.tensor(label)
        item['index'] = self.all_infos[idx]['index']

        return item

    def __len__(self):
        return len(self.all_infos)


# SMILES
class MolDatasetV2(torch.utils.data.Dataset):
    def __init__(self, dataset_path, encodings):
        self.encodings = encodings
        with open(dataset_path, 'r', encoding='UTF-8') as f:
            self.all_infos = json.load(f)

    def __getitem__(self, idx):
        item = {key: torch.tensor(val[idx])
                for key, val in self.encodings.items()}
        if self.all_infos[idx]['label'] == '-1':
            label = 0
        else:
            label = 1
        item['labels'] = torch.tensor(label)

        return item

    def __len__(self):
        return len(self.all_infos)