sample code

dummyData.py

import torch
from torch.utils.data import Dataset
import numpy as np

class DummyDataset(Dataset):
    def __init__(self, **kwargs):
        self.prob = kwargs['prob']
        self.vocab_size = kwargs['vocab_size']
        self.nSamples = kwargs['nSamples']
        self.src_data = np.random.choice(self.vocab_size, 
                                         self.nSamples)
    def __len__(self):
        return self.nSamples
    
    def func1(self, x):
        return x//2
    
    def func2(self, x):
        return 2*x + 1
    
    def get_target(self, x):
        if np.random.random() > self.prob:
            return self.func1(x)
        return self.func2(x)
    
    def __getitem__(self, index):
        assert index < self.nSamples
        x = self.src_data[index]
        y = self.get_target(x)
        return {'src':x, 'tgt':y}