| 562 | |
| 563 | |
| 564 | class MMapIndexedDatasetBuilder(object): |
| 565 | def __init__(self, out_file, dtype=np.int64): |
| 566 | self._data_file = open(out_file, "wb") |
| 567 | self._dtype = dtype |
| 568 | self._sizes = [] |
| 569 | self._doc_idx = [0] |
| 570 | |
| 571 | def add_item(self, tensor): |
| 572 | np_array = np.array(tensor.numpy(), dtype=self._dtype) |
| 573 | self._data_file.write(np_array.tobytes(order="C")) |
| 574 | self._sizes.append(np_array.size) |
| 575 | |
| 576 | def end_document(self): |
| 577 | self._doc_idx.append(len(self._sizes)) |
| 578 | |
| 579 | def merge_file_(self, another_file): |
| 580 | # Concatenate index |
| 581 | index = MMapIndexedDataset.Index(index_file_path(another_file)) |
| 582 | assert index.dtype == self._dtype |
| 583 | |
| 584 | for size in index.sizes: |
| 585 | self._sizes.append(size) |
| 586 | |
| 587 | # Concatenate data |
| 588 | with open(data_file_path(another_file), "rb") as f: |
| 589 | shutil.copyfileobj(f, self._data_file) |
| 590 | |
| 591 | def finalize(self, index_file): |
| 592 | self._data_file.close() |
| 593 | |
| 594 | with MMapIndexedDataset.Index.writer(index_file, self._dtype) as index: |
| 595 | index.write(self._sizes, self._doc_idx) |
no outgoing calls
no test coverage detected