රෙට්රෝපුහුණු දත්ත කට්ටලය

15import json
16from pathlib import Path
17
18import numpy as np
19import torch
20from torch.utils.data import Dataset as PyTorchDataset
21
22from labml import lab, monit
23from labml_helpers.datasets.text import TextFileDataset, TextDataset
24from labml_nn.transformers.retro.database import RetroIndex

27def build_dataset(chunk_len: int = 16, chunks_per_sample: int = 32, skip_range: int = 8):

39    dataset = TextFileDataset(
40        lab.get_data_path() / 'tiny_shakespeare.txt',
41        list,
42        url='https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt')

45    text = dataset.train

48    index = RetroIndex()

51    sample_offsets = []

53    i = 0
54    while i < len(text):

56        skip = np.random.randint(skip_range)
57        i += skip

60        if i + chunks_per_sample * chunk_len > len(text):
61            break

64        sample_offsets.append(i)

67        i += chunks_per_sample * chunk_len

70    samples = []

72    for i in monit.iterate('Gather Neighbors', sample_offsets):

74        sample = text[i: i + chunks_per_sample * chunk_len + 1]

76        src = sample[:-1]

78        chunks = [src[j:j + chunk_len] for j in range(0, len(src), chunk_len)]

80        chunk_offsets = [j + i for j in range(0, len(src), chunk_len)]

83        neighbor_offsets = index(chunks, chunk_offsets)

86        neighbors = [[text[j: j + chunk_len * 2] for j in n_off] for n_off in neighbor_offsets]

89        samples.append((sample[:-1], sample[1:], neighbors))

94    with open(str(lab.get_data_path() / 'retro_train_dataset.json'), 'w') as f:
95        f.write(json.dumps(samples))

98class Dataset(PyTorchDataset):

105    def __init__(self, file_path: Path, tds: TextDataset):

111        self.tds = tds

113        with open(str(file_path), 'r') as f:
114            self.samples = json.loads(f.read())

116    def __len__(self):

120        return len(self.samples)

122    def __getitem__(self, idx: int):

127        s = self.samples[idx]

129        src = self.tds.text_to_i(s[0])
130        tgt = self.tds.text_to_i(s[1])
131        neighbors = torch.stack([torch.stack([self.tds.text_to_i(n) for n in chunks]) for chunks in s[2]])

133        return src, tgt, neighbors

136if __name__ == '__main__':
137    build_dataset()