Coverage for /home/jenkins/.local/lib/python3.10/site-packages/hyper_parallel/data/indexed/indexed_pretrain

3# Licensed under the Apache License, Version 2.0 (the "License");

4# you may not use this file except in compliance with the License.

5# You may obtain a copy of the License at

7# http://www.apache.org/licenses/LICENSE-2.0

9# Unless required by applicable law or agreed to in writing, software

10# distributed under the License is distributed on an "AS IS" BASIS,

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

12# See the License for the specific language governing permissions and

13# limitations under the License.

14# ============================================================================

15"""Runtime Dataset types selected by the indexed pretraining builder."""

17from __future__ import annotations

19import hashlib

20import json

21import os

22import time

23from abc import ABC, abstractmethod

24from collections import OrderedDict

25from collections.abc import Mapping

26from typing import Any

28import numpy as np

29from numpy.random import RandomState

31from hyper_parallel.data.constants import IGNORE_INDEX, PAD_TOKEN_ID

32from hyper_parallel.data.dataset_logging import get_dataset_logger

33from hyper_parallel.data.indexed.indexed_data_config import GPTDatasetConfig

34from hyper_parallel.data.indexed.indexed_data_reader import ArrowDataReader, IndexedDataReader

36logger = get_dataset_logger(__name__)

39def _build_document_index(

40 documents: np.ndarray,

41 num_epochs: int,

42 random_state: RandomState,

43 separate_final_epoch: bool,

44) -> np.ndarray:

45 """Repeat and shuffle the exposed sequence indices for every epoch."""

46 if not separate_final_epoch or num_epochs == 1:

47 document_index = np.mgrid[0:num_epochs, 0:len(documents)][1]

48 document_index[:] = documents

49 document_index = document_index.reshape(-1).astype(np.int32)

50 random_state.shuffle(document_index)

51 return document_index

53 first_epochs = _build_document_index(documents, num_epochs - 1, random_state, False)

54 final_epoch = _build_document_index(documents, 1, random_state, False)

55 document_index = np.concatenate((first_epochs, final_epoch))

56 return document_index

59def _build_shuffle_index(

60 num_samples: int,

61 total_size: int,

62 random_state: RandomState,

63) -> np.ndarray:

64 """Shuffle the main epoch range separately from a small final epoch."""

65 dtype = np.uint32 if total_size < np.iinfo(np.uint32).max - 1 else np.int64

66 first_range = np.arange(0, num_samples, dtype=dtype)

67 random_state.shuffle(first_range)

68 if num_samples == total_size:

69 return first_range

71 final_range = np.arange(num_samples, total_size, dtype=dtype)

72 random_state.shuffle(final_range)

73 shuffle_index = np.concatenate((first_range, final_range))

74 return shuffle_index

77class _IndexedPretrainDataset(ABC):

78 """Store the common state used by all indexed pretraining Datasets."""

80 def __init__(

81 self,

82 dataset: IndexedDataReader | None,

83 dataset_path: str | None,

84 indices: np.ndarray | None,

85 num_samples: int,

86 index_split: Any,

87 config: GPTDatasetConfig,

88 ) -> None:

89 """Store common construction inputs and finalize the concrete Dataset."""

90 self.dataset = dataset

91 self.dataset_path = dataset_path

92 self.indices = indices

93 self.num_samples = num_samples

94 self.index_split = index_split

95 self.config = config

96 self._ltor_fields_cacheable = not any(

97 (

98 self.config.reset_position_ids,

99 self.config.reset_attention_mask,

100 self.config.eod_mask_loss,

101 )

102 )

103 self._cached_ltor_fields = None

104

105 # Handle pad token id provided by the tokenizer.

106 try:

107 pad_token_id = self.config.tokenizer.pad

108 self._pad_token_id = PAD_TOKEN_ID if pad_token_id is None else int(pad_token_id)

109 except (AttributeError, NotImplementedError):

110 self._pad_token_id = PAD_TOKEN_ID

111

112 try:

113 special_token_ids = [

114 token_id

115 for token_name, token_id in self.config.tokenizer.special_tokens_dict.items()

116 if token_name != "pad_token"

117 ]

118 except (AttributeError, IndexError, ValueError):

119 special_token_ids = []

120 if not special_token_ids:

121 for token_name in ("eos", "eod"):

122 try:

123 special_token_ids.append(getattr(self.config.tokenizer, token_name))

124 except (AttributeError, NotImplementedError):

125 pass

126

127 if not self.config.mock and self._pad_token_id in special_token_ids:

128 if self.config.allow_ambiguous_pad_tokens:

129 logger.warning(

130 "The tokenizer PAD ID collides with another special token ID; matching tokens will also be "

131 "masked as padding"

132 )

133 else:

134 self._pad_token_id = PAD_TOKEN_ID

135 logger.warning(

136 "The tokenizer PAD ID collides with another special token ID; using an internal padding sentinel"

137 )

138

139 if not self.config.mock:

140 self.unique_identifiers = type(self).build_unique_identifiers(

141 self.dataset_path,

142 self.num_samples,

143 self.index_split,

144 self.config,

145 )

146

147 self.unique_description = json.dumps(

148 self.unique_identifiers,

149 indent=4,

150 default=lambda value: value.unique_identifiers,

151 )

152 self.unique_description_hash = hashlib.md5(

153 self.unique_description.encode("utf-8")

154 ).hexdigest()

155

156 self._finalize()

157

158 def _build_ltor_fields(

159 self,

160 tokens: np.ndarray,

161 labels: np.ndarray,

162 ) -> dict[str, np.ndarray]:

163 """Build LTR masks and position IDs when explicitly enabled."""

164 if not self.config.create_ltor_fields:

165 ltor_fields = {}

166 return ltor_fields

167

168 if not self._ltor_fields_cacheable or self._cached_ltor_fields is None:

169 attention_mask, loss_mask, position_ids = _get_ltor_masks_and_position_ids(tokens, self.config)

170 if self._ltor_fields_cacheable:

171 self._cached_ltor_fields = (attention_mask, loss_mask, position_ids)

172 else:

173 attention_mask, loss_mask, position_ids = self._cached_ltor_fields

174

175 loss_mask = loss_mask.copy()

176 loss_mask[labels == self._pad_token_id] = 0.0

177 ltor_fields = {

178 "loss_mask": loss_mask,

179 "position_ids": position_ids,

180 "attention_mask": attention_mask,

181 }

182

183 return ltor_fields

184

185 def _finalize(self) -> None:

186 """Build indices or validate the low-level Dataset."""

187

188 @staticmethod

189 def numel_low_level_dataset(low_level_dataset: IndexedDataReader) -> int:

190 """Return the number of elements used to construct split indices."""

191 del low_level_dataset

192 raise NotImplementedError

193

194 @staticmethod

195 def build_low_level_dataset(

196 dataset_path: str,

197 config: GPTDatasetConfig,

198 ) -> IndexedDataReader:

199 """Build the low-level Dataset used by a concrete pretraining Dataset."""

200 del dataset_path, config

201 raise NotImplementedError

202

203 @staticmethod

204 def _key_config_attributes() -> list[str]:

205 """Return config fields that contribute to the Dataset cache identity."""

206 return [

207 "random_seed",

208 "sequence_length",

209 "split",

210 "split_matrix",

211 "tokenizer",

212 "drop_last_partial_validation_sequence",

213 "add_extra_token_to_sequence",

214 ]

215

216 @classmethod

217 def build_unique_identifiers(

218 cls,

219 dataset_path: str | None,

220 num_samples: int,

221 index_split: Any,

222 config: GPTDatasetConfig,

223 ) -> OrderedDict[str, Any]:

224 """Build the stable cache identity without constructing a Dataset.

225

226 Args:

227 dataset_path: Prefix of the indexed Dataset files.

228 num_samples: Requested number of samples.

229 index_split: Train, validation, or test split identifier.

230 config: Indexed GPT Dataset configuration.

231

232 Returns:

233 Ordered cache identity matching an instantiated Dataset.

234 """

235 identifiers = OrderedDict()

236 identifiers["class"] = cls.__name__

237 identifiers["dataset_path"] = dataset_path

238 identifiers["num_samples"] = num_samples

239 identifiers["index_split"] = index_split.name.lower()

240 for attribute in cls._key_config_attributes():

241 identifiers[attribute] = getattr(config, attribute)

242 return identifiers

243

244 @abstractmethod

245 def __len__(self) -> int:

246 """Return the number of samples exposed by the Dataset."""

247

248 @abstractmethod

249 def __getitem__(self, index: int) -> Mapping[str, Any]:

250 """Return one raw pretraining sample."""

251

252

253class _MockDataset(_IndexedPretrainDataset):

254 """Base class for mock Datasets that do not use a low-level reader."""

255

256 def __init__(

257 self,

258 dataset: IndexedDataReader | None,

259 dataset_path: str | None,

260 indices: np.ndarray | None,

261 num_samples: int,

262 index_split: Any,

263 config: GPTDatasetConfig,

264 ) -> None:

265 """Validate mock mode and initialize the common Dataset state."""

266 self.config = config

267 if not self.config.mock:

268 raise ValueError("config.mock must be True for MockDataset")

269 super().__init__(dataset, dataset_path, indices, num_samples, index_split, config)

270

271 def __len__(self) -> int:

272 """Return the configured number of mock samples."""

273 return self.num_samples

274

275

276class MockGPTDataset(_MockDataset):

277 """Generate deterministic, fixed-length mock GPT samples."""

278

279 def __getitem__(self, index: int) -> Mapping[str, Any]:

280 """Generate a deterministic mock sample for one split and index."""

281 token = 1

282 pad = 2

283 eod = self.config.tokenizer.eod

284

285 if index >= self.num_samples:

286 raise ValueError(

287 f"Exceeded the available number of samples ({self.num_samples}); got {index}"

288 )

289

290 rng = np.random.default_rng(seed=[self.index_split.value, index])

291 length = rng.integers(low=0, high=self.config.sequence_length)

292 sample_tokens = np.zeros(length) + token

293 sample_pads = np.zeros(self.config.sequence_length - length - 1) + pad

294 sample = np.int64(np.concatenate([[length], sample_tokens, [eod], sample_pads]))

295 tokens = sample[:-1]

296 labels = sample[1:]

297

298 raw_sample = {

299 "input_ids": tokens,

300 "labels": labels,

301 }

302 raw_sample.update(self._build_ltor_fields(tokens, labels))

303

304 return raw_sample

305

306

307class GPTDataset(_IndexedPretrainDataset):

308 """Build fixed-length GPT samples from an indexed token stream."""

309

310 def _finalize(self) -> None:

311 """Build the document, sample, and shuffle indices for the token stream."""

312 if self.dataset is None or self.indices is None:

313 raise ValueError("GPTDataset requires a low-level Dataset and split indices")

314

315 self.document_index, self.sample_index, self.shuffle_index = self._build_document_sample_shuffle_indices()

316

317 def _plan_epochs(self, num_tokens_per_epoch: int) -> tuple[int, bool, int]:

318 """Choose the epoch count and whether to shuffle a short final epoch separately."""

319 num_epochs = 0

320 accumulated_tokens = 0

321 extra_token = int(self.config.add_extra_token_to_sequence)

322 requested_tokens = self.num_samples * self.config.sequence_length + extra_token

323 while accumulated_tokens < requested_tokens:

324 num_epochs += 1

325 accumulated_tokens += num_tokens_per_epoch

326

327 if num_epochs == 1:

328 separate_final_epoch = False

329 samples_without_final_epoch = 0

330 else:

331 samples_without_final_epoch = (

332 (num_epochs - 1) * num_tokens_per_epoch - extra_token

333 ) // self.config.sequence_length

334 samples_from_final_epoch = self.num_samples - samples_without_final_epoch

335 samples_per_epoch = (num_tokens_per_epoch - extra_token) // self.config.sequence_length

336

337 # Keep a small final epoch separate from the main shuffled range.

338 threshold = 0.80

339 separate_final_epoch = samples_from_final_epoch < int(threshold * samples_per_epoch)

340 return num_epochs, separate_final_epoch, samples_without_final_epoch

341

342 def _build_document_sample_shuffle_indices(self) -> tuple[np.ndarray, np.ndarray, np.ndarray]:

343 """Build or load the document, sample, and shuffle indices.

344

345 The document index:

346 -- 1-D

347 -- An ordered array of document ids

348

349 The sample index:

350 -- 2-D

351 -- The document indices and offsets which mark the start of every sample

352

353 The shuffle index:

354 -- 1-D

355 -- A random permutation of index range of the sample index

356

357 Returns:

358 The document, sample, and shuffle index arrays.

359 """

360 # pylint: disable=too-many-locals

361 cache_directory = self.config.path_to_cache

362 if cache_directory is None:

363 cache_directory = os.path.join(self.dataset.path_prefix, "cache", f"{type(self).__name__}_indices")

364 cache_key = f"{self.unique_description_hash}-{type(self).__name__}"

365

366 # drop last for train/test, valid drop depended on config

367 if self.index_split.name.lower() == "valid":

368 drop_last_partial_sequence = self.config.drop_last_partial_validation_sequence

369 else:

370 drop_last_partial_sequence = True

371

372 cache_paths = {

373 "description": os.path.join(cache_directory, f"{cache_key}-description.txt"),

374 "document": os.path.join(cache_directory, f"{cache_key}-document_index.npy"),

375 "sample": os.path.join(cache_directory, f"{cache_key}-sample_index.npy"),

376 "shuffle": os.path.join(cache_directory, f"{cache_key}-shuffle_index.npy"),

377 }

378

379 # read from cache

380 cache_hit = all(os.path.isfile(path) for path in cache_paths.values())

381 if cache_hit:

382 logger.debug("Loading indexed sample cache: key=%s, directory=%s", cache_key, cache_directory, enabled=True)

383 document_index = np.load(cache_paths["document"], mmap_mode="r")

384 sample_index = np.load(cache_paths["sample"], mmap_mode="r")

385 shuffle_index = np.load(cache_paths["shuffle"], mmap_mode="r")

386

387 return document_index, sample_index, shuffle_index

388

389 # build cache

390 start_time = time.time()

391 num_tokens_per_epoch = int(np.sum(self.dataset.sequence_lengths[self.indices]))

392 if num_tokens_per_epoch <= 1:

393 raise ValueError("The selected indexed split must contain at least two tokens")

394

395 num_epochs, separate_final_epoch, samples_without_final_epoch = self._plan_epochs(num_tokens_per_epoch)

396 random_state = RandomState(self.config.random_seed)

397

398 # Shuffle the repeated document IDs, optionally isolating the final epoch.

399 document_index = _build_document_index(self.indices, num_epochs, random_state, separate_final_epoch)

400

401 # Convert the shuffled token stream into fixed-length sample boundaries in C++.

402 # Keep the optional C++ extension out of the pre-packed MindRecord path.

403 from hyper_parallel.data.indexed.indexed_helpers import build_sample_index # pylint: disable=C0415

404

405 sample_index = build_sample_index(

406 self.dataset.sequence_lengths,

407 document_index,

408 self.config.sequence_length,

409 num_epochs,

410 num_tokens_per_epoch,

411 drop_last_partial_sequence,

412 self.config.add_extra_token_to_sequence,

413 )

414

415 # Shuffle samples without mixing an isolated final epoch into the main range.

416 if separate_final_epoch:

417 shuffle_index = _build_shuffle_index(samples_without_final_epoch, sample_index.shape[0] - 1, random_state)

418 else:

419 shuffle_index = _build_shuffle_index(sample_index.shape[0] - 1, sample_index.shape[0] - 1, random_state)

420

421 os.makedirs(cache_directory, exist_ok=True)

422 np.save(cache_paths["document"], document_index)

423 np.save(cache_paths["sample"], sample_index)

424 np.save(cache_paths["shuffle"], shuffle_index)

425 # Write the description last so readers treat it as the cache completion marker.

426 with open(cache_paths["description"], "w", encoding="utf-8") as description_file:

427 description_file.write(self.unique_description)

428

429 logger.debug(

430 "Built indexed sample cache: samples=%d, epochs=%d, elapsed=%.4f seconds",

431 self.num_samples, num_epochs, time.time() - start_time,

432 )

433 return document_index, sample_index, shuffle_index

434

435 @staticmethod

436 def numel_low_level_dataset(low_level_dataset: IndexedDataReader) -> int:

437 """Return the number of low-level sequences available for splitting."""

438 return int(low_level_dataset.sequence_lengths.shape[0])

439

440 @staticmethod

441 def build_low_level_dataset(

442 dataset_path: str,

443 config: GPTDatasetConfig,

444 ) -> IndexedDataReader:

445 """Open the low-level indexed token reader."""

446 return IndexedDataReader(dataset_path, mmap=config.mmap_bin_files)

447

448 def __len__(self) -> int:

449 """Return the number of complete fixed-length samples."""

450 return int(self.sample_index.shape[0] - 1)

451

452 def __getitem__(self, index: int) -> Mapping[str, Any]:

453 """Read, shift, and mask one shuffled GPT sample."""

454 text, _, _ = self._query_document_sample_shuffle_indices(index)

455

456 if self.config.add_extra_token_to_sequence:

457 tokens = text[:-1].copy()

458 labels = text[1:].copy()

459 else:

460 tokens = text.copy()

461 labels = np.roll(text, shift=-1)

462 labels[-1] = self._pad_token_id

463

464 ltor_fields = self._build_ltor_fields(tokens, labels)

465

466 # For padded sequences, ensure the embedding layer can map the token ID

467 tokens[tokens == self._pad_token_id] = 0

468 labels[labels == self._pad_token_id] = 0

469

470 if np.any(tokens < 0) or np.any(tokens >= len(self.config.tokenizer)):

471 raise ValueError("An input token is out of bounds of the tokenizer vocabulary")

472

473 raw_sample = {

474 "input_ids": tokens,

475 "labels": labels,

476 }

477 raw_sample.update(ltor_fields)

478

479 return raw_sample

480

481 def _query_document_sample_shuffle_indices(

482 self,

483 index: int,

484 ) -> tuple[np.ndarray, np.ndarray]:

485 """Get the text (token ids) and document ids for a given index

486

487 Args:

488 index (int): The index into the dataset

489

490 Returns:

491 The text ids, document ids, and per-document token counts (before any padding).

492 """

493 shuffled_index = int(self.shuffle_index[index])

494 begin_document, begin_offset = self.sample_index[shuffled_index]

495 end_document, end_offset = self.sample_index[shuffled_index + 1]

496

497 document_ids = []

498 sample_parts = []

499

500 # Read one contiguous slice when the sample stays within a document.

501 if begin_document == end_document:

502 document_id = int(self.document_index[begin_document])

503 document_ids.append(document_id)

504 sample_parts.append(

505 self.dataset.get(

506 document_id,

507 offset=int(begin_offset),

508 length=int(end_offset - begin_offset + self.config.add_extra_token_to_sequence),

509 )

510 )

511

512 # Concatenate slices when the sample crosses document boundaries.

513 else:

514 for document_position in range(int(begin_document), int(end_document) + 1):

515 document_id = int(self.document_index[document_position])

516 document_ids.append(document_id)

517 offset = 0 if document_position > begin_document else int(begin_offset)

518 length = (

519 None

520 if document_position < end_document

521 else int(end_offset + self.config.add_extra_token_to_sequence)

522 )

523 sample_parts.append(self.dataset.get(document_id, offset=offset, length=length))

524

525 # Retained partial samples still need the configured fixed length.

526 length = sum(map(len, sample_parts))

527 document_lengths = [len(p) for p in sample_parts]

528

529 if length < (self.config.sequence_length + self.config.add_extra_token_to_sequence):

530 sample_parts.append(

531 [self._pad_token_id]

532 * (self.config.sequence_length + self.config.add_extra_token_to_sequence - length)

533 )

534

535 text = np.asarray(np.concatenate(sample_parts), dtype=np.int64)

536

537 return text, np.asarray(document_ids, dtype=np.int64), document_lengths

538

539

540class GPTFromMRDataset(_IndexedPretrainDataset):

541 """GPT Dataset for pre-packed samples converted from MindRecord.

542

543 Each entry in the underlying indexed Dataset represents a complete pre-packed

544 training sample. Samples are read directly without rebuilding document, sample,

545 or shuffle indices.

546

547 Args:

548 dataset: The indexed Dataset containing pre-packed samples converted from MindRecord.

549

550 dataset_path: The path to the underlying indexed Dataset.

551

552 indices: The indices of the pre-packed samples exposed by this Dataset.

553

554 num_samples: The number of pre-packed samples to draw from the indexed Dataset.

555

556 index_split: The Dataset split associated with ``indices``.

557

558 config: The GPT-specific Dataset configuration.

559 """

560

561 def _finalize(self) -> None:

562 """Validate the low-level Dataset and the pre-cut record lengths."""

563 if self.dataset is None or self.indices is None:

564 raise ValueError("GPTFromMRDataset requires a low-level Dataset and split indices")

565

566 if not isinstance(self.config, GPTDatasetConfig):

567 raise ValueError("GPTFromMRDataset requires GPTDatasetConfig")

568

569 if not self.config.skip_data_check:

570 expected_length = self.config.sequence_length + self.config.add_extra_token_to_sequence

571 if not np.all(self.dataset.sequence_lengths == expected_length):

572 raise ValueError(

573 "All pre-cut records must match sequence_length + add_extra_token_to_sequence; "

574 f"expected {expected_length}, got {self.dataset.sequence_lengths.tolist()}"

575 )

576 if self.num_samples > len(self.indices):

577 raise ValueError(

578 "Requested samples exceed the pre-cut records in this split: "

579 f"requested {self.num_samples}, available {len(self.indices)}"

580 )

581

582 @staticmethod

583 def numel_low_level_dataset(low_level_dataset: IndexedDataReader) -> int:

584 """Return the number of pre-cut records available for splitting."""

585 return int(low_level_dataset.sequence_lengths.shape[0])

586

587 @staticmethod

588 def is_multimodal() -> bool:

589 """Return whether records contain multimodal sequence modes."""

590 return False

591

592 @staticmethod

593 def is_split_by_sequence() -> bool:

594 """Return whether train/validation/test boundaries use sequence indices."""

595 return True

596

597 @staticmethod

598 def build_low_level_dataset(

599 dataset_path: str,

600 config: GPTDatasetConfig,

601 ) -> IndexedDataReader:

602 """Open a low-level reader for pre-cut records."""

603 return IndexedDataReader(

604 dataset_path,

605 mmap=config.mmap_bin_files,

606 reuse_index=config.reuse_idx,

607 )

608

609 def __len__(self) -> int:

610 """Return the requested number of pre-cut records in this split."""

611 return self.num_samples

612

613 def __getitem__(self, index: int) -> Mapping[str, Any]:

614 """Read and shift one pre-cut record without dynamic concatenation."""

615 sequence_id = int(self.indices[index])

616 text = np.asarray(self.dataset[sequence_id], dtype=np.int64)

617 if np.any(text == self._pad_token_id):

618 raise ValueError("Pre-cut indexed Dataset records must be packed without padding")

619

620 if self.config.add_extra_token_to_sequence:

621 tokens = text[:-1].copy()

622 labels = text[1:].copy()

623 else:

624 tokens = text.copy()

625 labels = np.roll(text, shift=-1)

626 labels[-1] = self._pad_token_id

627

628 if np.any(tokens < 0) or np.any(tokens >= self.config.tokenizer.vocab_size):

629 raise ValueError("An input token is out of bounds of the tokenizer vocabulary")

630

631 ltor_fields = self._build_ltor_fields(tokens, labels)

632

633 tokens[tokens == self._pad_token_id] = 0

634 labels[labels == self._pad_token_id] = 0

635

636 raw_sample = {

637 "input_ids": tokens,

638 "labels": labels,

639 }

640 raw_sample.update(ltor_fields)

641

642 return raw_sample

643

644

645def _get_ltor_masks_and_position_ids(

646 tokens: np.ndarray,

647 config: GPTDatasetConfig,

648) -> tuple[np.ndarray, np.ndarray, np.ndarray]:

649 """Build causal attention, loss, and position arrays for one token sequence.

650

651 Args:

652 tokens: Token IDs for one fixed-length Indexed sample.

653 config: Indexed Dataset runtime-field configuration.

654

655 Returns:

656 Attention mask, loss mask, and position IDs.

657 """

658 seq_length = tokens.size

659 eod_token = config.tokenizer.eod

660 attention_mask = np.tril(np.ones((seq_length, seq_length)))[None, :, :]

661

662 # Loss mask.

663 loss_mask = np.ones(seq_length, dtype=np.float32)

664 if config.eod_mask_loss:

665 loss_mask[tokens == eod_token] = 0.0

666

667 # Position ids.

668 position_ids = np.arange(seq_length, dtype=np.int64)

669 # We need to clone as the ids will be modified based on batch index.

670 if config.reset_position_ids:

671 position_ids = position_ids.copy()

672

673 if config.reset_position_ids or config.reset_attention_mask:

674 # Find indices where EOD token is.

675 eod_indices = position_ids[tokens == eod_token]

676 # Detach indices from positions if going to modify positions.

677 if config.reset_position_ids:

678 eod_indices = eod_indices.copy()

679

680 # Loop through EOD indices:

681 prev_index = 0

682 for eod_index in eod_indices:

683 next_index = int(eod_index) + 1

684 # Mask attention loss.

685 if config.reset_attention_mask:

686 attention_mask[0, next_index:, :next_index] = 0

687 # Reset positions.

688 if config.reset_position_ids:

689 position_ids[next_index:] -= next_index - prev_index

690 prev_index = next_index

691

692 # Convert attention mask to binary:

693 attention_mask = attention_mask < 0.5

694

695 return attention_mask, loss_mask, position_ids

696

697

698class SFTFromMRDataset:

699 """Prepared Arrow records with authoritative next-token labels and boundaries.

700

701 Each row is kept intact, without tokenization, shifting, EOS insertion or

702 truncation. Field names can be mapped inside ``low_level_dataset`` using

703 input_ids_key, labels_key, boundaries_key and loss_mask_key. Missing boundaries

704 mean one document. Optional binary loss_mask is folded into owned labels.

705 """

706

707 def __init__(

708 self, dataset: ArrowDataReader, dataset_path: str, indices: np.ndarray,

709 num_samples: int, index_split: Any, config: GPTDatasetConfig,

710 ) -> None:

711 """Store whole-record sampling inputs without building GPT packing caches."""

712 self.dataset = dataset

713 self.indices = indices

714 self.num_samples = num_samples

715 self.config = config

716 self.unique_identifiers = self.build_unique_identifiers(dataset_path, num_samples, index_split, config)

717 if self.config.simple_blend not in ("inter", "intra"):

718 raise ValueError("Prepared Arrow datasets require simple_blend='inter' or 'intra'")

719 if self.num_samples > len(self.indices):

720 raise ValueError("Requested samples exceed prepared records in this split")

721

722 @staticmethod

723 def numel_low_level_dataset(low_level_dataset: ArrowDataReader) -> int:

724 """Split by complete records rather than by inner documents."""

725 return len(low_level_dataset)

726

727 @staticmethod

728 def build_low_level_dataset(dataset_path: str, config: GPTDatasetConfig) -> ArrowDataReader:

729 """Open the Arrow reader used by the existing split builder."""

730 del config

731 return ArrowDataReader(dataset_path)

732

733 @classmethod

734 def build_unique_identifiers(

735 cls, dataset_path: str, num_samples: int, index_split: Any, config: GPTDatasetConfig,

736 ) -> dict[str, Any]:

737 """Identify the record view without opening Arrow files during lazy setup."""

738 return {"class": cls.__name__, "dataset_path": dataset_path, "num_samples": num_samples,

739 "index_split": index_split.name.lower(), "split": config.split,

740 "dataset_options": config.dataset_options}

741

742 def __len__(self) -> int:

743 """Return the requested record count."""

744 return self.num_samples

745

746 def __getitem__(self, index: int) -> Mapping[str, Any]:

747 """Read labels unchanged except for explicit binary mask suppression."""

748 sample = self.dataset[int(self.indices[index])]

749 fields = self.config.dataset_options

750 tokens = np.asarray(sample[fields.get("input_ids_key", "input_ids")])

751 labels = np.asarray(sample[fields.get("labels_key", "labels")])

752 boundaries = np.asarray(sample.get(fields.get("boundaries_key", "cu_seq_lens"), [0, tokens.size]))

753 boundaries = self._validate_record_layout(tokens, labels, boundaries, int(self.indices[index]))

754 if (np.any(tokens < 0) or np.any(tokens >= self.config.tokenizer.vocab_size)

755 or np.any((labels < 0) & (labels != IGNORE_INDEX))

756 or np.any(labels >= self.config.tokenizer.vocab_size)):

757 raise ValueError("Prepared token/label IDs must fit the vocabulary; ignored labels must be -100")

758 # Arrow buffers may be shared/read-only; never modify their supervision in place.

759 labels = labels.astype(np.int64, copy=True)

760 mask = sample.get(fields.get("loss_mask_key", "loss_mask"))

761 if mask is not None:

762 mask = np.asarray(mask)

763 if mask.shape != labels.shape or not np.isin(mask, (0, 1)).all():

764 raise ValueError("loss_mask must be binary and match labels")

765 labels[mask == 0] = IGNORE_INDEX

766 if not np.any(labels != IGNORE_INDEX):

767 raise ValueError("Prepared record has no effective targets")

768 return {"input_ids": tokens.astype(np.int64), "labels": labels,

769 "cu_seq_lens": boundaries.astype(np.int32)}

770

771 def _validate_record_layout(

772 self, tokens: np.ndarray, labels: np.ndarray, boundaries: np.ndarray, record_index: int,

773 ) -> np.ndarray:

774 """Check record shapes, configured length and internal document boundaries."""

775 if (tokens.ndim != 1 or not tokens.size or labels.shape != tokens.shape

776 or tokens.dtype.kind not in "iu" or labels.dtype.kind not in "iu"):

777 raise ValueError("Prepared tokens and labels must be matching nonempty 1D integer arrays")

778 if not self.config.skip_data_check and tokens.size != self.config.sequence_length:

779 raise ValueError(

780 "Prepared record must match seq_length; "

781 f"record {record_index}: expected {self.config.sequence_length}, got {tokens.size}"

782 )

783 boundaries = boundaries[boundaries != -1] if boundaries.ndim == 1 else boundaries

784 if (boundaries.ndim != 1 or boundaries.dtype.kind not in "iu" or boundaries.size < 2

785 or boundaries[0] != 0 or boundaries[-1] != tokens.size

786 or np.any(boundaries[1:] <= boundaries[:-1]) or tokens.size > np.iinfo(np.int32).max):

787 raise ValueError("Document boundaries must strictly increase from zero to token count")

788 return boundaries

Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / data / indexed / indexed_pretrain_dataset.py: 79%

360 statements