(self, data_root, tokenizer, label_name, train_length=4096, min_length=512, num_data=-1, seed=42, dataset_ckpt_path=None)
| 17 | class StreamingTrainingJsonlZSD(torch.utils.data.Dataset): |
| 18 | |
| 19 | def __init__(self, data_root, tokenizer, label_name, train_length=4096, min_length=512, num_data=-1, seed=42, dataset_ckpt_path=None): |
| 20 | |
| 21 | self.data_root = data_root |
| 22 | |
| 23 | self.data_path = sorted([f'{data_root}/{path}' for path in os.listdir(data_root) if os.path.isdir(f'{data_root}/{path}') and 'git' not in path]) |
| 24 | self.data_path = sorted(sum([[f'{data_root}/{path}' for path in os.listdir(data_root)] for data_root in self.data_path], [])) |
| 25 | self.data_path = sorted(sum([[f'{data_root}/{path}' for path in os.listdir(data_root)] for data_root in self.data_path], [])) |
| 26 | |
| 27 | random.shuffle(self.data_path) |
| 28 | self.dctx = zstandard.ZstdDecompressor() |
| 29 | |
| 30 | self.tokenizer = tokenizer |
| 31 | self.label_name = label_name |
| 32 | |
| 33 | self.len = num_data |
| 34 | self.train_length = train_length |
| 35 | self.min_length = min_length |
| 36 | |
| 37 | self.pivot = torch.distributed.get_rank() |
| 38 | self.size = torch.distributed.get_world_size() |
| 39 | |
| 40 | self.token_buffer, self.file_buffer = [], None |
| 41 | |
| 42 | self.file_buffer = open(self.data_path[self.pivot], 'rb') |
| 43 | self.file_buffer = self.dctx.stream_reader(self.file_buffer) |
| 44 | self.file_buffer = io.TextIOWrapper(self.file_buffer, encoding='utf-8') |
| 45 | self.sample_idx = 0 |
| 46 | |
| 47 | if dataset_ckpt_path is not None: |
| 48 | dataset_ckpt_path = f"{dataset_ckpt_path}/dataset_ckpt-{self.pivot:{len(str(self.size))}d}-{self.size}.pt" |
| 49 | dataset_ckpt = torch.load(dataset_ckpt_path, weights_only=False) |
| 50 | self.data_path = dataset_ckpt['data_path'] |
| 51 | self.label_name = dataset_ckpt['label_name'] |
| 52 | self.pivot = dataset_ckpt['pivot'] |
| 53 | self.size = dataset_ckpt['size'] |
| 54 | self.file_buffer = open(self.data_path[self.pivot], 'rb') |
| 55 | self.file_buffer = self.dctx.stream_reader(self.file_buffer) |
| 56 | self.file_buffer = io.TextIOWrapper(self.file_buffer, encoding='utf-8') |
| 57 | self.sample_idx = dataset_ckpt['sample_idx'] |
| 58 | for _ in range(self.sample_idx): |
| 59 | sample = self.file_buffer.readline() |
| 60 | self.token_buffer = dataset_ckpt['token_buffer'] |
| 61 | |
| 62 | def __len__(self): |
| 63 | return self.len |
nothing calls this directly
no outgoing calls
no test coverage detected