amachine.am_transformers.am_datasets
1import json 2import logging 3from typing import Iterator 4import numpy as np 5import pyarrow.parquet as pq 6from torch.utils.data import IterableDataset 7import torch 8import warnings 9from tokenizers import Tokenizer 10from collections import deque 11from torch.multiprocessing import Value 12 13logger = logging.getLogger(__name__) 14 15class StreamingParquetDataset(IterableDataset): 16 17 18 def __init__( 19 self, 20 path: str, 21 input_column: str, 22 remap: np.ndarray, 23 seq_len: int, 24 shuffle_buffer_size : int = 1024, 25 additional_columns: tuple[str, ...] = (), 26 rename_input_ids : bool = True, 27 stride: int | None = None, 28 shuffle_windows: bool = True, 29 preserve_file_order : bool = False, 30 rank: int = 0, 31 world_size: int = 1, 32 seed : int = 32 33 ): 34 35 super().__init__() 36 37 if preserve_file_order : 38 39 if shuffle_windows : 40 raise ValueError( f"shuffle_windows not compatable with preserve_file_order" ) 41 42 if world_size > 1: 43 raise ValueError( 44 "preserve_global_order=True requires world_size=1. " 45 "Strict global order cannot be maintained across distributed ranks." 46 ) 47 48 elif shuffle_windows and shuffle_buffer_size <= 2 : 49 raise ValueError( f"shuffle_buffer_size should be greater than 2" ) 50 51 if rename_input_ids and input_column != "input_ids": 52 if "input_ids" in additional_columns: 53 raise ValueError( 54 "Conflict: 'rename_input_ids' is True, but 'input_ids' is already " 55 "present in 'additional_columns'. This will cause data to be overwritten." 56 ) 57 58 self.path = path 59 self.input_col = input_column 60 61 seen = set() 62 columns = [] 63 for col in [input_column, *additional_columns]: 64 if col in seen: 65 warnings.warn(f"Duplicate column {col!r} ignored", stacklevel=2) 66 else: 67 seen.add(col) 68 columns.append(col) 69 70 self.columns : list[str] = columns 71 72 self.preserve_file_order = preserve_file_order 73 self.remap = remap 74 self.seq_len = seq_len 75 self.shuffle_buffer_size = shuffle_buffer_size 76 self.shuffle_windows = shuffle_windows 77 self._shared_epoch = Value('i', 0) 78 self.rename_input_ids = rename_input_ids 79 self.rank = rank 80 self.world_size = world_size 81 self.target_len = seq_len + 1 82 self.stride = self.target_len - 1 if stride is None else stride 83 self.seed = seed 84 85 if not (1 <= self.stride <= self.target_len) : 86 raise ValueError(f"stride must be in [1, {self.target_len}], got {self.stride}") 87 88 def _get_global_worker_info(self): 89 """Calculates global worker ID based on explicit constructor args.""" 90 91 worker_info = torch.utils.data.get_worker_info() 92 worker_id = worker_info.id if worker_info is not None else 0 93 num_workers = worker_info.num_workers if worker_info is not None else 1 94 95 if self.preserve_file_order and (self.world_size > 1 or num_workers > 1): 96 raise ValueError( 97 "preserve_file_order=True requires world_size=1 and num_workers=1" 98 ) 99 100 global_num_workers = self.world_size * num_workers 101 global_worker_id = (self.rank * num_workers) + worker_id 102 103 return global_worker_id, global_num_workers 104 105 def set_epoch(self, epoch: int) -> None: 106 if epoch < 0: 107 raise ValueError("epoch must be non-negative") 108 109 self._shared_epoch.value = epoch 110 111 def _make_rng( 112 self, 113 global_worker_id: int, 114 epoch: int, 115 ) -> np.random.Generator: 116 117 seed_sequence = np.random.SeedSequence( 118 self.seed, 119 spawn_key=(epoch, global_worker_id ), 120 ) 121 122 return np.random.default_rng(seed_sequence) 123 124 def _row_group_indices( 125 self, 126 global_worker_id: int, 127 global_num_workers : int, 128 n_rg: int) -> range: 129 130 if global_num_workers == 1: 131 return range(n_rg) 132 133 groups_per_worker, remainder = divmod(n_rg, global_num_workers ) 134 135 start = ( 136 global_worker_id * groups_per_worker 137 + min( global_worker_id, remainder ) 138 ) 139 140 stop = ( 141 start 142 + groups_per_worker 143 + int( global_worker_id < remainder ) 144 ) 145 146 return range(start, stop) 147 148 def _read_row_group( 149 self, 150 parquet: pq.ParquetFile, 151 rg_idx: int, 152 ) -> dict[str, np.ndarray] | None : 153 154 table = parquet.read_row_group(rg_idx, columns=self.columns) 155 row = {} 156 157 for col in self.columns: 158 column = table.column(col) 159 160 if column.null_count: 161 raise ValueError( 162 f"Column {col} contains nulls in row group {rg_idx}" 163 ) 164 165 values = column.to_numpy(zero_copy_only=False) 166 167 if not np.issubdtype(values.dtype, np.integer): 168 raise TypeError( 169 f"Column {col!r} must be integer-valued, " 170 f"got {values.dtype} in row group {rg_idx}" 171 ) 172 173 row[col] = values.astype(np.int64, copy=False) 174 175 arr = row[self.input_col] 176 if arr.size == 0: 177 return None 178 179 lo = int(arr.min()) 180 hi = int(arr.max()) 181 182 if lo < 0 or hi >= len(self.remap): 183 raise ValueError( 184 f"Tokens out of remap range in row group {rg_idx}: " 185 f"observed [{lo}, {hi}], expected [0, {len(self.remap) - 1}]" 186 ) 187 188 row[self.input_col] = self.remap[arr] 189 return row 190 191 def _stream_windows( 192 self, 193 parquet: pq.ParquetFile, 194 rg_indices: range, 195 skip: int = 0, 196 ) -> Iterator[dict[str, np.ndarray]]: 197 """ 198 Concatenate row groups and yield fixed-length windows in order. 199 `skip` drops that many leading tokens (used for the per-epoch offset). 200 Trailing tokens that don't fill a window are discarded. 201 """ 202 target_len, stride = self.target_len, self.stride 203 pool = {col: np.empty(0, dtype=np.int64) for col in self.columns} 204 205 for rg_idx in rg_indices: 206 207 row = self._read_row_group(parquet, rg_idx) 208 209 if row is None: 210 continue 211 212 # Could be done more efficiently since concatenate requires reallocation 213 for col in self.columns: 214 pool[col] = np.concatenate([pool[col], row[col]]) 215 216 # Consume any remaining skip before cutting windows. 217 if skip: 218 n = len(pool[self.input_col]) 219 if n <= skip: 220 skip -= n 221 pool = {col: a[:0] for col, a in pool.items()} 222 continue 223 pool = {col: a[skip:] for col, a in pool.items()} 224 skip = 0 225 226 n = len(pool[self.input_col]) 227 228 consumed = 0 229 for pos in range(0, n - target_len + 1, stride): 230 yield {col: pool[col][pos: pos + target_len].copy() 231 for col in self.columns} 232 consumed = pos + stride 233 234 if consumed: 235 pool = {col: pool[col][consumed:].copy() for col in self.columns} 236 237 def _iter_shuffled_windows( 238 self, 239 parquet : pq.ParquetFile, 240 rg_indices, 241 rng, 242 epoch ) : 243 244 skip = 0 if epoch == 0 else int(rng.integers(0, self.stride)) 245 buffer: deque[dict[str, np.ndarray]] = deque() 246 247 def emit(): 248 idx = int(rng.integers(len(buffer))) 249 buffer[idx], buffer[-1] = buffer[-1], buffer[idx] 250 return buffer.pop() 251 252 for window in self._stream_windows(parquet, rg_indices, skip=skip): 253 buffer.append(window) 254 while len(buffer) >= self.shuffle_buffer_size: 255 yield emit() 256 257 while buffer: 258 yield emit() 259 260 def _iter_windows(self) -> Iterator[dict[str, np.ndarray]]: 261 262 with pq.ParquetFile(self.path) as parquet: 263 264 epoch = self._shared_epoch.value 265 266 global_worker_id, global_n_workers = self._get_global_worker_info() 267 268 rg_indices = self._row_group_indices( 269 global_worker_id, 270 global_n_workers, 271 parquet.metadata.num_row_groups 272 ) 273 274 if len(rg_indices) == 0: 275 return 276 277 if self.shuffle_windows: 278 rng = self._make_rng( global_worker_id, epoch ) 279 yield from self._iter_shuffled_windows(parquet, rg_indices, rng, epoch) 280 else: 281 yield from self._stream_windows(parquet, rg_indices) 282 283 def __iter__(self) -> Iterator[dict[str, torch.Tensor]]: 284 285 for window in self._iter_windows(): 286 287 item = { 288 col: torch.from_numpy(values) 289 for col, values in window.items() 290 } 291 292 if self.rename_input_ids and self.input_col != "input_ids" : 293 item["input_ids"] = item.pop(self.input_col) 294 295 item["attention_mask"] = torch.ones_like(item["input_ids"], dtype=torch.long) 296 297 yield item 298 299def build_remap_table( 300 metadata_path: str, 301 tokenizer : Tokenizer, 302 unk_token : str | None = "<|unk|>", 303 strict: bool = True 304) -> np.ndarray: 305 306 with open(metadata_path, 'r', encoding='utf-8') as f: 307 meta = json.load(f) 308 309 alphabet: list[str] = meta["alphabet"] 310 vocab = tokenizer.get_vocab() 311 unk_id = tokenizer.token_to_id( unk_token ) 312 313 if unk_id is None : 314 raise ValueError(f"Tokenizer missing {unk_token} token.") 315 316 table = np.full(len(alphabet), fill_value=unk_id, dtype=np.int64) 317 318 missing = [] 319 for i, ch in enumerate(alphabet): 320 tid = vocab.get(ch) 321 if tid is None: 322 missing.append(ch) 323 logger.debug("Character %r not in vocab -> UNK (%d)", ch, unk_id) 324 else: 325 table[i] = tid 326 logger.debug("Mapped %r -> %d", ch, tid) 327 328 if strict and missing: 329 raise ValueError(f"Strict mode: alphabet has unmapped chars: {missing}") 330 331 return table
logger =
<Logger amachine.am_transformers.am_datasets (WARNING)>
class
StreamingParquetDataset(torch.utils.data.dataset.Dataset[+_T_co], typing.Iterable[+_T_co]):
16class StreamingParquetDataset(IterableDataset): 17 18 19 def __init__( 20 self, 21 path: str, 22 input_column: str, 23 remap: np.ndarray, 24 seq_len: int, 25 shuffle_buffer_size : int = 1024, 26 additional_columns: tuple[str, ...] = (), 27 rename_input_ids : bool = True, 28 stride: int | None = None, 29 shuffle_windows: bool = True, 30 preserve_file_order : bool = False, 31 rank: int = 0, 32 world_size: int = 1, 33 seed : int = 32 34 ): 35 36 super().__init__() 37 38 if preserve_file_order : 39 40 if shuffle_windows : 41 raise ValueError( f"shuffle_windows not compatable with preserve_file_order" ) 42 43 if world_size > 1: 44 raise ValueError( 45 "preserve_global_order=True requires world_size=1. " 46 "Strict global order cannot be maintained across distributed ranks." 47 ) 48 49 elif shuffle_windows and shuffle_buffer_size <= 2 : 50 raise ValueError( f"shuffle_buffer_size should be greater than 2" ) 51 52 if rename_input_ids and input_column != "input_ids": 53 if "input_ids" in additional_columns: 54 raise ValueError( 55 "Conflict: 'rename_input_ids' is True, but 'input_ids' is already " 56 "present in 'additional_columns'. This will cause data to be overwritten." 57 ) 58 59 self.path = path 60 self.input_col = input_column 61 62 seen = set() 63 columns = [] 64 for col in [input_column, *additional_columns]: 65 if col in seen: 66 warnings.warn(f"Duplicate column {col!r} ignored", stacklevel=2) 67 else: 68 seen.add(col) 69 columns.append(col) 70 71 self.columns : list[str] = columns 72 73 self.preserve_file_order = preserve_file_order 74 self.remap = remap 75 self.seq_len = seq_len 76 self.shuffle_buffer_size = shuffle_buffer_size 77 self.shuffle_windows = shuffle_windows 78 self._shared_epoch = Value('i', 0) 79 self.rename_input_ids = rename_input_ids 80 self.rank = rank 81 self.world_size = world_size 82 self.target_len = seq_len + 1 83 self.stride = self.target_len - 1 if stride is None else stride 84 self.seed = seed 85 86 if not (1 <= self.stride <= self.target_len) : 87 raise ValueError(f"stride must be in [1, {self.target_len}], got {self.stride}") 88 89 def _get_global_worker_info(self): 90 """Calculates global worker ID based on explicit constructor args.""" 91 92 worker_info = torch.utils.data.get_worker_info() 93 worker_id = worker_info.id if worker_info is not None else 0 94 num_workers = worker_info.num_workers if worker_info is not None else 1 95 96 if self.preserve_file_order and (self.world_size > 1 or num_workers > 1): 97 raise ValueError( 98 "preserve_file_order=True requires world_size=1 and num_workers=1" 99 ) 100 101 global_num_workers = self.world_size * num_workers 102 global_worker_id = (self.rank * num_workers) + worker_id 103 104 return global_worker_id, global_num_workers 105 106 def set_epoch(self, epoch: int) -> None: 107 if epoch < 0: 108 raise ValueError("epoch must be non-negative") 109 110 self._shared_epoch.value = epoch 111 112 def _make_rng( 113 self, 114 global_worker_id: int, 115 epoch: int, 116 ) -> np.random.Generator: 117 118 seed_sequence = np.random.SeedSequence( 119 self.seed, 120 spawn_key=(epoch, global_worker_id ), 121 ) 122 123 return np.random.default_rng(seed_sequence) 124 125 def _row_group_indices( 126 self, 127 global_worker_id: int, 128 global_num_workers : int, 129 n_rg: int) -> range: 130 131 if global_num_workers == 1: 132 return range(n_rg) 133 134 groups_per_worker, remainder = divmod(n_rg, global_num_workers ) 135 136 start = ( 137 global_worker_id * groups_per_worker 138 + min( global_worker_id, remainder ) 139 ) 140 141 stop = ( 142 start 143 + groups_per_worker 144 + int( global_worker_id < remainder ) 145 ) 146 147 return range(start, stop) 148 149 def _read_row_group( 150 self, 151 parquet: pq.ParquetFile, 152 rg_idx: int, 153 ) -> dict[str, np.ndarray] | None : 154 155 table = parquet.read_row_group(rg_idx, columns=self.columns) 156 row = {} 157 158 for col in self.columns: 159 column = table.column(col) 160 161 if column.null_count: 162 raise ValueError( 163 f"Column {col} contains nulls in row group {rg_idx}" 164 ) 165 166 values = column.to_numpy(zero_copy_only=False) 167 168 if not np.issubdtype(values.dtype, np.integer): 169 raise TypeError( 170 f"Column {col!r} must be integer-valued, " 171 f"got {values.dtype} in row group {rg_idx}" 172 ) 173 174 row[col] = values.astype(np.int64, copy=False) 175 176 arr = row[self.input_col] 177 if arr.size == 0: 178 return None 179 180 lo = int(arr.min()) 181 hi = int(arr.max()) 182 183 if lo < 0 or hi >= len(self.remap): 184 raise ValueError( 185 f"Tokens out of remap range in row group {rg_idx}: " 186 f"observed [{lo}, {hi}], expected [0, {len(self.remap) - 1}]" 187 ) 188 189 row[self.input_col] = self.remap[arr] 190 return row 191 192 def _stream_windows( 193 self, 194 parquet: pq.ParquetFile, 195 rg_indices: range, 196 skip: int = 0, 197 ) -> Iterator[dict[str, np.ndarray]]: 198 """ 199 Concatenate row groups and yield fixed-length windows in order. 200 `skip` drops that many leading tokens (used for the per-epoch offset). 201 Trailing tokens that don't fill a window are discarded. 202 """ 203 target_len, stride = self.target_len, self.stride 204 pool = {col: np.empty(0, dtype=np.int64) for col in self.columns} 205 206 for rg_idx in rg_indices: 207 208 row = self._read_row_group(parquet, rg_idx) 209 210 if row is None: 211 continue 212 213 # Could be done more efficiently since concatenate requires reallocation 214 for col in self.columns: 215 pool[col] = np.concatenate([pool[col], row[col]]) 216 217 # Consume any remaining skip before cutting windows. 218 if skip: 219 n = len(pool[self.input_col]) 220 if n <= skip: 221 skip -= n 222 pool = {col: a[:0] for col, a in pool.items()} 223 continue 224 pool = {col: a[skip:] for col, a in pool.items()} 225 skip = 0 226 227 n = len(pool[self.input_col]) 228 229 consumed = 0 230 for pos in range(0, n - target_len + 1, stride): 231 yield {col: pool[col][pos: pos + target_len].copy() 232 for col in self.columns} 233 consumed = pos + stride 234 235 if consumed: 236 pool = {col: pool[col][consumed:].copy() for col in self.columns} 237 238 def _iter_shuffled_windows( 239 self, 240 parquet : pq.ParquetFile, 241 rg_indices, 242 rng, 243 epoch ) : 244 245 skip = 0 if epoch == 0 else int(rng.integers(0, self.stride)) 246 buffer: deque[dict[str, np.ndarray]] = deque() 247 248 def emit(): 249 idx = int(rng.integers(len(buffer))) 250 buffer[idx], buffer[-1] = buffer[-1], buffer[idx] 251 return buffer.pop() 252 253 for window in self._stream_windows(parquet, rg_indices, skip=skip): 254 buffer.append(window) 255 while len(buffer) >= self.shuffle_buffer_size: 256 yield emit() 257 258 while buffer: 259 yield emit() 260 261 def _iter_windows(self) -> Iterator[dict[str, np.ndarray]]: 262 263 with pq.ParquetFile(self.path) as parquet: 264 265 epoch = self._shared_epoch.value 266 267 global_worker_id, global_n_workers = self._get_global_worker_info() 268 269 rg_indices = self._row_group_indices( 270 global_worker_id, 271 global_n_workers, 272 parquet.metadata.num_row_groups 273 ) 274 275 if len(rg_indices) == 0: 276 return 277 278 if self.shuffle_windows: 279 rng = self._make_rng( global_worker_id, epoch ) 280 yield from self._iter_shuffled_windows(parquet, rg_indices, rng, epoch) 281 else: 282 yield from self._stream_windows(parquet, rg_indices) 283 284 def __iter__(self) -> Iterator[dict[str, torch.Tensor]]: 285 286 for window in self._iter_windows(): 287 288 item = { 289 col: torch.from_numpy(values) 290 for col, values in window.items() 291 } 292 293 if self.rename_input_ids and self.input_col != "input_ids" : 294 item["input_ids"] = item.pop(self.input_col) 295 296 item["attention_mask"] = torch.ones_like(item["input_ids"], dtype=torch.long) 297 298 yield item
An iterable Dataset.
All datasets that represent an iterable of data samples should subclass it. Such form of datasets is particularly useful when data come from a stream.
All subclasses should overwrite __iter__(), which would return an
iterator of samples in this dataset.
When a subclass is used with ~torch.utils.data.DataLoader, each
item in the dataset will be yielded from the ~torch.utils.data.DataLoader
iterator. When num_workers > 0, each worker process will have a
different copy of the dataset object, so it is often desired to configure
each copy independently to avoid having duplicate data returned from the
workers. ~torch.utils.data.get_worker_info(), when called in a worker
process, returns information about the worker. It can be used in either the
dataset's __iter__() method or the ~torch.utils.data.DataLoader 's
worker_init_fn option to modify each copy's behavior.
Example 1: splitting workload across all workers in __iter__()::
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_DATALOADER)
>>> # xdoctest: +SKIP("Fails on MacOS12")
>>> class MyIterableDataset(torch.utils.data.IterableDataset):
... def __init__(self, start, end):
... super(MyIterableDataset).__init__()
... assert end > start, "this example only works with end >= start"
... self.start = start
... self.end = end
...
... def __iter__(self):
... worker_info = torch.utils.data.get_worker_info()
... if worker_info is None: # single-process data loading, return the full iterator
... iter_start = self.start
... iter_end = self.end
... else: # in a worker process
... # split workload
... per_worker = int(math.ceil((self.end - self.start) / float(worker_info.num_workers)))
... worker_id = worker_info.id
... iter_start = self.start + worker_id * per_worker
... iter_end = min(iter_start + per_worker, self.end)
... return iter(range(iter_start, iter_end))
...
>>> # should give same set of data as range(3, 7), i.e., [3, 4, 5, 6].
>>> ds = MyIterableDataset(start=3, end=7)
>>> # Single-process loading
>>> print(list(torch.utils.data.DataLoader(ds, num_workers=0)))
[tensor([3]), tensor([4]), tensor([5]), tensor([6])]
>>> # xdoctest: +REQUIRES(POSIX)
>>> # Multi-process loading with two worker processes
>>> # Worker 0 fetched [3, 4]. Worker 1 fetched [5, 6].
>>> # xdoctest: +IGNORE_WANT("non deterministic")
>>> print(list(torch.utils.data.DataLoader(ds, num_workers=2)))
[tensor([3]), tensor([5]), tensor([4]), tensor([6])]
>>> # With even more workers
>>> # xdoctest: +IGNORE_WANT("non deterministic")
>>> print(list(torch.utils.data.DataLoader(ds, num_workers=12)))
[tensor([3]), tensor([5]), tensor([4]), tensor([6])]
Example 2: splitting workload across all workers using worker_init_fn::
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_DATALOADER)
>>> class MyIterableDataset(torch.utils.data.IterableDataset):
... def __init__(self, start, end):
... super(MyIterableDataset).__init__()
... assert end > start, "this example only works with end >= start"
... self.start = start
... self.end = end
...
... def __iter__(self):
... return iter(range(self.start, self.end))
...
>>> # should give same set of data as range(3, 7), i.e., [3, 4, 5, 6].
>>> ds = MyIterableDataset(start=3, end=7)
>>> # Single-process loading
>>> print(list(torch.utils.data.DataLoader(ds, num_workers=0)))
[3, 4, 5, 6]
>>>
>>> # Directly doing multi-process loading yields duplicate data
>>> print(list(torch.utils.data.DataLoader(ds, num_workers=2)))
[3, 3, 4, 4, 5, 5, 6, 6]
>>> # Define a `worker_init_fn` that configures each dataset copy differently
>>> def worker_init_fn(worker_id):
... worker_info = torch.utils.data.get_worker_info()
... dataset = worker_info.dataset # the dataset copy in this worker process
... overall_start = dataset.start
... overall_end = dataset.end
... # configure the dataset to only process the split workload
... per_worker = int(math.ceil((overall_end - overall_start) / float(worker_info.num_workers)))
... worker_id = worker_info.id
... dataset.start = overall_start + worker_id * per_worker
... dataset.end = min(dataset.start + per_worker, overall_end)
...
>>> # Mult-process loading with the custom `worker_init_fn`
>>> # Worker 0 fetched [3, 4]. Worker 1 fetched [5, 6].
>>> print(list(torch.utils.data.DataLoader(ds, num_workers=2, worker_init_fn=worker_init_fn)))
[3, 5, 4, 6]
>>> # With even more workers
>>> print(list(torch.utils.data.DataLoader(ds, num_workers=12, worker_init_fn=worker_init_fn)))
[3, 4, 5, 6]
StreamingParquetDataset( path: str, input_column: str, remap: numpy.ndarray, seq_len: int, shuffle_buffer_size: int = 1024, additional_columns: tuple[str, ...] = (), rename_input_ids: bool = True, stride: int | None = None, shuffle_windows: bool = True, preserve_file_order: bool = False, rank: int = 0, world_size: int = 1, seed: int = 32)
19 def __init__( 20 self, 21 path: str, 22 input_column: str, 23 remap: np.ndarray, 24 seq_len: int, 25 shuffle_buffer_size : int = 1024, 26 additional_columns: tuple[str, ...] = (), 27 rename_input_ids : bool = True, 28 stride: int | None = None, 29 shuffle_windows: bool = True, 30 preserve_file_order : bool = False, 31 rank: int = 0, 32 world_size: int = 1, 33 seed : int = 32 34 ): 35 36 super().__init__() 37 38 if preserve_file_order : 39 40 if shuffle_windows : 41 raise ValueError( f"shuffle_windows not compatable with preserve_file_order" ) 42 43 if world_size > 1: 44 raise ValueError( 45 "preserve_global_order=True requires world_size=1. " 46 "Strict global order cannot be maintained across distributed ranks." 47 ) 48 49 elif shuffle_windows and shuffle_buffer_size <= 2 : 50 raise ValueError( f"shuffle_buffer_size should be greater than 2" ) 51 52 if rename_input_ids and input_column != "input_ids": 53 if "input_ids" in additional_columns: 54 raise ValueError( 55 "Conflict: 'rename_input_ids' is True, but 'input_ids' is already " 56 "present in 'additional_columns'. This will cause data to be overwritten." 57 ) 58 59 self.path = path 60 self.input_col = input_column 61 62 seen = set() 63 columns = [] 64 for col in [input_column, *additional_columns]: 65 if col in seen: 66 warnings.warn(f"Duplicate column {col!r} ignored", stacklevel=2) 67 else: 68 seen.add(col) 69 columns.append(col) 70 71 self.columns : list[str] = columns 72 73 self.preserve_file_order = preserve_file_order 74 self.remap = remap 75 self.seq_len = seq_len 76 self.shuffle_buffer_size = shuffle_buffer_size 77 self.shuffle_windows = shuffle_windows 78 self._shared_epoch = Value('i', 0) 79 self.rename_input_ids = rename_input_ids 80 self.rank = rank 81 self.world_size = world_size 82 self.target_len = seq_len + 1 83 self.stride = self.target_len - 1 if stride is None else stride 84 self.seed = seed 85 86 if not (1 <= self.stride <= self.target_len) : 87 raise ValueError(f"stride must be in [1, {self.target_len}], got {self.stride}")
def
build_remap_table( metadata_path: str, tokenizer: tokenizers.Tokenizer, unk_token: str | None = '<|unk|>', strict: bool = True) -> numpy.ndarray:
300def build_remap_table( 301 metadata_path: str, 302 tokenizer : Tokenizer, 303 unk_token : str | None = "<|unk|>", 304 strict: bool = True 305) -> np.ndarray: 306 307 with open(metadata_path, 'r', encoding='utf-8') as f: 308 meta = json.load(f) 309 310 alphabet: list[str] = meta["alphabet"] 311 vocab = tokenizer.get_vocab() 312 unk_id = tokenizer.token_to_id( unk_token ) 313 314 if unk_id is None : 315 raise ValueError(f"Tokenizer missing {unk_token} token.") 316 317 table = np.full(len(alphabet), fill_value=unk_id, dtype=np.int64) 318 319 missing = [] 320 for i, ch in enumerate(alphabet): 321 tid = vocab.get(ch) 322 if tid is None: 323 missing.append(ch) 324 logger.debug("Character %r not in vocab -> UNK (%d)", ch, unk_id) 325 else: 326 table[i] = tid 327 logger.debug("Mapped %r -> %d", ch, tid) 328 329 if strict and missing: 330 raise ValueError(f"Strict mode: alphabet has unmapped chars: {missing}") 331 332 return table