Class TextDataset

text_loader.py:8–23

class TextDataset(Dataset):
    # Initialize your data, download, etc.

    def __init__(self, filename="./data/shakespeare.txt.gz"):
        self.len = 0
        with gzip.open(filename, 'rt') as f:
            self.targetLines = [x.strip() for x in f if x.strip()]
            self.srcLines = [x.lower().replace(' ', '')
                             for x in self.targetLines]
            self.len = len(self.srcLines)

    def __getitem__(self, index):
        return self.srcLines[index], self.targetLines[index]

    def __len__(self):
        return self.len


# Test the loader