Class MultiTaskDataset

preprocess/to_model.py:133–148

class MultiTaskDataset(Dataset):

    def __init__(self, meta_args, name2dataset: Dict[str, Dataset], split: str):
        self.meta_args = meta_args

        self.data = MultiTaskWrapper(name2dataset=name2dataset, meta_args=meta_args, split=split)

    def __getitem__(self, index):
        data = self.data[index]
        model_inputs = {
            k: data[k] for k in data['model_kwargs']
        }
        return model_inputs

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