chatbot article part 2_a

2_a.py

class DataGenerator():
    def __init__(self, dset):
        self.dset = dset
        self.len = len(self.dset)
        self.idx = 0
    def __len__(self):
        return self.len
    def __iter__(self):
        return self
    def __next__(self):
        x = Variable(torch.LongTensor([self.dset[self.idx:self.idx+2]]))
        if self.idx == self.len - 2:
            raise StopIteration
        y = Variable(torch.LongTensor([self.dset[self.idx+2]]), requires_grad=False)
        self.idx = self.idx + 1
        return (x, y)