Create a dataset class

Changes

1 changed files (+26/-0)

• dataset.py (new)
  • @@ -0,0 +1,26 @@
    + import torch + + + class Dataset(torch.utils.data.Dataset): + def __init__(self, file, seq_size): + self.seq_size = 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] + + 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]))