TransFew/MyDataset.py at main · BioinfoMachineLearning/TransFew · 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
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
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
import torch
from torch.utils.data import Dataset
from sklearn.utils.class_weight import compute_class_weight
import numpy as np
from Utils import pickle_load
import CONSTANTS

class TransFewDataset(Dataset):
    def __init__(self, data_pth=None,  submodel=None):

        self.submodel = submodel

        data = pickle_load(data_pth)

        labels = data['labels']
        labels = torch.cat(labels, dim=0)

        self.labels = labels
        self.esm_features = data['esm2_t48']
        self.msa_features = data['msa_1b']
        #self.diamond_features = data['diamond']
        self.interpro_features = data['interpro']
        #self.string_features = data['string']


    def __getitem__(self, index):

        esm = self.esm_features[index]
        msa = self.msa_features[index]
        # diamond = self.diamond_features[index]
        interpro = self.interpro_features[index]
        # string = self.string_features[index]
        label = self.labels[index]

        if self.submodel == 'esm2_t48':
            return esm, label
        elif self.submodel == 'msa_1b':
            return msa, label
        elif self.submodel == 'interpro':
            return interpro, label
        elif self.submodel == 'full':
            # return esm, msa, diamond, interpro, string, label
            return esm, msa, interpro, label

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


class TestDataset(Dataset):
    def __init__(self, data_pth=None, submodel=None):

        self.submodel = submodel

        data = pickle_load(data_pth)

        self.proteins = data['protein']
        self.esm_features = data['esm2_t48']
        self.msa_features = data['msa_1b']
        self.diamond_features = data['diamond']
        self.interpro_features = data['interpro']
        # self.labs = data['labels']

    def __getitem__(self, index):

        esm = self.esm_features[index]
        msa = self.msa_features[index]
        diamond = self.diamond_features[index]
        interpro = self.interpro_features[index]
        proteins = self.proteins[index]

        # pop = self.labs[index]


        if self.submodel == 'esm2_t48':
            return esm, proteins
        elif self.submodel == 'msa_1b':
            return msa, proteins
        elif self.submodel == 'diamond':
            return diamond, proteins
        elif self.submodel == 'interpro':
            return interpro, proteins
        elif self.submodel == 'full':
            return esm, msa, diamond, interpro, proteins#, pop

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


class PredictDataset(Dataset):
    def __init__(self, data=None):

        self.esm_features = data['esm2_t48']
        self.proteins = data['protein']


    def __getitem__(self, index):

        esm = self.esm_features[index]
        proteins = self.proteins[index]

        return esm, proteins

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