| 543 | |
| 544 | |
| 545 | class JsonVisionDataset(object): |
| 546 | @staticmethod |
| 547 | def get_default_config(updates=None): |
| 548 | config = ConfigDict() |
| 549 | config.path = '' |
| 550 | config.seq_length = 384 |
| 551 | config.batch_size = 4 |
| 552 | config.always_start_with_bos = False |
| 553 | config.start_seek_loc = 0 |
| 554 | config.example_index_at_start = 0 |
| 555 | config.tokens_count_at_start = 0 |
| 556 | config.tokenizer_processes = 1 |
| 557 | config.tokenizer_parallel_chunk_size = 32 |
| 558 | config.tokenizer_parallel_batch_size = 1024 |
| 559 | config.throughput_average_window_size = 200 |
| 560 | config.use_data_sharded_loader = True |
| 561 | config.return_local_batch = False |
| 562 | config.mode = 'pad' |
| 563 | |
| 564 | if updates is not None: |
| 565 | config.update(ConfigDict(updates).copy_and_resolve_references()) |
| 566 | return config |
| 567 | |
| 568 | def __init__(self, config, tokenizer, text_processor, node_info): |
| 569 | self.config = self.get_default_config(config) |
| 570 | assert self.config.path != '' |
| 571 | self._node_info = node_info |
| 572 | self._tokenizer = tokenizer |
| 573 | self._text_processor = text_processor |
| 574 | self._index = self.config.example_index_at_start |
| 575 | self._file_loc = self.config.start_seek_loc |
| 576 | self._total_tokens = 0 |
| 577 | |
| 578 | def parse_json(self, line): |
| 579 | if not line or line == '\n': |
| 580 | return None |
| 581 | try: |
| 582 | data = json.loads(line) |
| 583 | except json.decoder.JSONDecodeError: |
| 584 | print(f'Error parsing json line:\n{line}') |
| 585 | return None |
| 586 | return data |
| 587 | |
| 588 | def json_iterator(self): |
| 589 | index, file_loc = self._index, self._file_loc |
| 590 | with open_file(self.config.path, 'r', block_size=50 * 2 ** 20) as fin: |
| 591 | fin.seek(file_loc) |
| 592 | while True: |
| 593 | line = fin.readline() |
| 594 | file_loc = fin.tell() |
| 595 | if not line: # Reached EOF |
| 596 | index = 0 |
| 597 | fin.seek(0) |
| 598 | continue |
| 599 | if not self.config.use_data_sharded_loader or index % self._node_info['dp_node_size'] == self._node_info['dp_node_rank']: |
| 600 | data = self.parse_json(line) |
| 601 | if data is not None: |
| 602 | # JSON parsing succeeded |