1. 1
2. 2
3. 3
4. 4
5. 5
6. 6
7. 7
8. 8
9. 9
10. 10
11. 11
12. 12
13. 13
14. 14
15. 15
16. 16
17. 17
18. 18
19. 19
20. 20
21. 21
22. 22
23. 23
24. 24
25. 25
26. 26
27. 27
28. 28

from collections import Counter

import torch


class Dataset(torch.utils.data.Dataset):
    def __init__(self, file, seq_size):
        with open(file, 'r') as f:
            self.words = f.read().split()

        self.word_counts = Counter(self.words)
        self.uniq_words = sorted(self.word_counts, key=self.word_counts.get)

        self.index_to_word = {index: word for index,
                              word in enumerate(self.uniq_words)}
        self.word_to_index = {word: index for index,
                              word in enumerate(self.uniq_words)}

        self.words_indexes = [self.word_to_index[w] for w in self.words]

        self.seq_size = seq_size

    def __len__(self):
        return len(self.words_indexes) - self.seq_size

    def __getitem__(self, index):
        return (torch.tensor(self.words_indexes[index:index+self.seq_size]),
                torch.tensor(self.words_indexes[index+1:index+self.seq_size+1]))