A class designed to manage the packing and storage of large arrays into binary files with a specified chunk size. This class handles the division of large arrays into smaller 'chunks' that are stored individually in binary files. Each file begins with a header specifying metadata such
| 79 | |
| 80 | |
| 81 | class PackedDatasetBuilder(object): |
| 82 | """ |
| 83 | A class designed to manage the packing and storage of large arrays into binary files with a specified chunk size. |
| 84 | |
| 85 | This class handles the division of large arrays into smaller 'chunks' that are stored individually in binary files. |
| 86 | Each file begins with a header specifying metadata such as data type and version. This facilitates the management |
| 87 | of potentially large datasets that need to be processed or transmitted in smaller, more manageable units. |
| 88 | |
| 89 | Packing data works like this: |
| 90 | |
| 91 | 1) A big array of chunk size is created with prefilled with pad tokens. |
| 92 | 2) When #add_array is called and given the tokenized |
| 93 | |
| 94 | Parameters: |
| 95 | outdir (str): The output directory where the chunk files will be stored. |
| 96 | prefix (str): The prefix to use for naming the chunk files. |
| 97 | chunk_size (int): The maximum number of elements each chunk file should contain. |
| 98 | pad_token (int): Incomplete chunks will be filled with pad_token. |
| 99 | dtype (str or numpy.dtype, optional): The data type of the array elements. If 'auto', the dtype is determined based on `vocab_size`. |
| 100 | Defaults to 'auto'. |
| 101 | vocab_size (int, optional): The maximum size of the vocabulary. Required if dtype is 'auto'. |
| 102 | """ |
| 103 | |
| 104 | def __init__( |
| 105 | self, |
| 106 | outdir: StrPath, |
| 107 | prefix: str, |
| 108 | chunk_size: int, |
| 109 | pad_token: int, |
| 110 | dtype='auto', |
| 111 | vocab_size=None, |
| 112 | ): |
| 113 | if dtype == 'auto': |
| 114 | if vocab_size is None: |
| 115 | raise ValueError("vocab_size cannot be None when dtype='auto'") |
| 116 | if vocab_size is not None and vocab_size < 65500: |
| 117 | self._dtype = np.uint16 |
| 118 | else: |
| 119 | self._dtype = np.int32 |
| 120 | else: |
| 121 | self._dtype = dtype |
| 122 | self._counter = 0 |
| 123 | self._chunk_size = chunk_size |
| 124 | self._outdir = outdir |
| 125 | self._prefix = prefix |
| 126 | self._pad_token = pad_token |
| 127 | |
| 128 | # Initialise an array with the pad tokens to fill up as we turn file contents into tokens |
| 129 | self._arr = np.zeros(self._chunk_size, dtype=self._dtype) |
| 130 | self._arr.fill(self._pad_token) |
| 131 | |
| 132 | self._idx = 0 |
| 133 | self._version = 1 |
| 134 | self._filenames = [] |
| 135 | |
| 136 | def _write_chunk(self): |
| 137 | filename = f'{self._prefix}_{self._counter:010d}.bin' |
| 138 | filename = os.path.join(self._outdir, filename) |
no outgoing calls