from EIVData.csv_dataset import CSVData from torch.utils.data import random_split def load_data(seed=0, splitting_part=0.8, normalize=True): """ Loads the protein structure dataset :param seed: Seed for splitting and shuffling the data. Defaults to 0. :param splitting_part: Which fraction of the data to use as training data. Defaults to 0.8. :normalize: Whether to normalize the data, defaults to True. :returns: protein_trainset, protein_testset """ protein_dataset = CSVData('~/SharedData/AI/datasets/protein_structure/CASP.csv', class_name='RMSD', shuffle_seed=seed, normalize=normalize, delimiter=",") dataset_len = len(protein_dataset) train_len = int(dataset_len*splitting_part) test_len = dataset_len - train_len protein_trainset, protein_testset = random_split(protein_dataset, lengths=[train_len, test_len]) return protein_trainset, protein_testset