from torch.utils.data import Dataset


class MalwareDatasetForNN(Dataset):
    """malware pytorch dataset for neural nets"""

    def __init__(self, data, nn_conf):
        self.data = data
        self.nn_conf = nn_conf

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

    def __getitem__(self, index):
        data = self.data[index]
        
        sample = dict()
        sample['label'] = data.HasDetections
        for c in self.nn_conf['cat'].keys():
            sample[c] = data[c]
        sample['cont'] = data[self.nn_conf['cont']]

        return sample
