Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / data / indexed / indexed_pretrain_dataset.py: 0%
306 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
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
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
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 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 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
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
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
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 )
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 )
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()
156 self._finalize()
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
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
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 }
183 return ltor_fields
185 def _finalize(self) -> None:
186 """Build indices or validate the low-level Dataset."""
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
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
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 ]
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.
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.
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
244 @abstractmethod
245 def __len__(self) -> int:
246 """Return the number of samples exposed by the Dataset."""
248 @abstractmethod
249 def __getitem__(self, index: int) -> Mapping[str, Any]:
250 """Return one raw pretraining sample."""
253class _MockDataset(_IndexedPretrainDataset):
254 """Base class for mock Datasets that do not use a low-level reader."""
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)
271 def __len__(self) -> int:
272 """Return the configured number of mock samples."""
273 return self.num_samples
276class MockGPTDataset(_MockDataset):
277 """Generate deterministic, fixed-length mock GPT samples."""
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
285 if index >= self.num_samples:
286 raise ValueError(
287 f"Exceeded the available number of samples ({self.num_samples}); got {index}"
288 )
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:]
298 raw_sample = {
299 "tokens": tokens,
300 "labels": labels,
301 }
302 raw_sample.update(self._build_ltor_fields(tokens, labels))
304 return raw_sample
307class GPTDataset(_IndexedPretrainDataset):
308 """Build fixed-length GPT samples from an indexed token stream."""
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")
315 self.document_index, self.sample_index, self.shuffle_index = self._build_document_sample_shuffle_indices()
317 def _build_document_sample_shuffle_indices(self) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
318 """Build or load the document, sample, and shuffle indices.
320 The document index:
321 -- 1-D
322 -- An ordered array of document ids
324 The sample index:
325 -- 2-D
326 -- The document indices and offsets which mark the start of every sample
328 The shuffle index:
329 -- 1-D
330 -- A random permutation of index range of the sample index
332 Returns:
333 The document, sample, and shuffle index arrays.
334 """
335 # pylint: disable=too-many-locals
336 cache_directory = self.config.path_to_cache
337 if cache_directory is None:
338 cache_directory = os.path.join(self.dataset.path_prefix, "cache", f"{type(self).__name__}_indices")
339 cache_key = f"{self.unique_description_hash}-{type(self).__name__}"
341 # drop last for train/test, valid drop depended on config
342 if self.index_split.name.lower() == "valid":
343 drop_last_partial_sequence = self.config.drop_last_partial_validation_sequence
344 else:
345 drop_last_partial_sequence = True
347 cache_paths = {
348 "description": os.path.join(cache_directory, f"{cache_key}-description.txt"),
349 "document": os.path.join(cache_directory, f"{cache_key}-document_index.npy"),
350 "sample": os.path.join(cache_directory, f"{cache_key}-sample_index.npy"),
351 "shuffle": os.path.join(cache_directory, f"{cache_key}-shuffle_index.npy"),
352 }
354 # read from cache
355 cache_hit = all(os.path.isfile(path) for path in cache_paths.values())
356 if cache_hit:
357 logger.debug("Loading indexed sample cache: key=%s, directory=%s", cache_key, cache_directory, enabled=True)
358 document_index = np.load(cache_paths["document"], mmap_mode="r")
359 sample_index = np.load(cache_paths["sample"], mmap_mode="r")
360 shuffle_index = np.load(cache_paths["shuffle"], mmap_mode="r")
362 return document_index, sample_index, shuffle_index
364 # build cache
365 start_time = time.time()
366 num_tokens_per_epoch = int(np.sum(self.dataset.sequence_lengths[self.indices]))
367 if num_tokens_per_epoch <= 1:
368 raise ValueError("The selected indexed split must contain at least two tokens")
370 # repeate documents for multi-epcho
371 num_epochs = 0
372 accumulated_tokens = 0
373 extra_token = int(self.config.add_extra_token_to_sequence)
374 requested_tokens = self.num_samples * self.config.sequence_length + extra_token
375 while accumulated_tokens < requested_tokens:
376 num_epochs += 1
377 accumulated_tokens += num_tokens_per_epoch
379 if num_epochs == 1:
380 separate_final_epoch = False
381 samples_without_final_epoch = 0
382 else:
383 samples_without_final_epoch = (
384 (num_epochs - 1) * num_tokens_per_epoch - extra_token
385 ) // self.config.sequence_length
386 samples_from_final_epoch = self.num_samples - samples_without_final_epoch
387 samples_per_epoch = (num_tokens_per_epoch - extra_token) // self.config.sequence_length
389 # Keep a small final epoch separate from the main shuffled range.
390 threshold = 0.80
391 separate_final_epoch = samples_from_final_epoch < int(threshold * samples_per_epoch)
393 random_state = RandomState(self.config.random_seed)
395 # Shuffle the repeated document IDs, optionally isolating the final epoch.
396 document_index = _build_document_index(self.indices, num_epochs, random_state, separate_final_epoch)
398 # Convert the shuffled token stream into fixed-length sample boundaries in C++.
399 # Keep the optional C++ extension out of the pre-packed MindRecord path.
400 from hyper_parallel.data.indexed.indexed_helpers import build_sample_index # pylint: disable=C0415
402 sample_index = build_sample_index(
403 self.dataset.sequence_lengths,
404 document_index,
405 self.config.sequence_length,
406 num_epochs,
407 num_tokens_per_epoch,
408 drop_last_partial_sequence,
409 self.config.add_extra_token_to_sequence,
410 )
412 # Shuffle samples without mixing an isolated final epoch into the main range.
413 if separate_final_epoch:
414 shuffle_index = _build_shuffle_index(samples_without_final_epoch, sample_index.shape[0] - 1, random_state)
415 else:
416 shuffle_index = _build_shuffle_index(sample_index.shape[0] - 1, sample_index.shape[0] - 1, random_state)
418 os.makedirs(cache_directory, exist_ok=True)
419 np.save(cache_paths["document"], document_index)
420 np.save(cache_paths["sample"], sample_index)
421 np.save(cache_paths["shuffle"], shuffle_index)
422 # Write the description last so readers treat it as the cache completion marker.
423 with open(cache_paths["description"], "w", encoding="utf-8") as description_file:
424 description_file.write(self.unique_description)
426 logger.debug(
427 "Built indexed sample cache: samples=%d, epochs=%d, elapsed=%.4f seconds",
428 self.num_samples, num_epochs, time.time() - start_time,
429 )
430 return document_index, sample_index, shuffle_index
432 @staticmethod
433 def numel_low_level_dataset(low_level_dataset: IndexedDataReader) -> int:
434 """Return the number of low-level sequences available for splitting."""
435 return int(low_level_dataset.sequence_lengths.shape[0])
437 @staticmethod
438 def build_low_level_dataset(
439 dataset_path: str,
440 config: GPTDatasetConfig,
441 ) -> IndexedDataReader:
442 """Open the low-level indexed token reader."""
443 return IndexedDataReader(dataset_path, mmap=config.mmap_bin_files)
445 def __len__(self) -> int:
446 """Return the number of complete fixed-length samples."""
447 return int(self.sample_index.shape[0] - 1)
449 def __getitem__(self, index: int) -> Mapping[str, Any]:
450 """Read, shift, and mask one shuffled GPT sample."""
451 text, _, _ = self._query_document_sample_shuffle_indices(index)
453 if self.config.add_extra_token_to_sequence:
454 tokens = text[:-1].copy()
455 labels = text[1:].copy()
456 else:
457 tokens = text.copy()
458 labels = np.roll(text, shift=-1)
459 labels[-1] = self._pad_token_id
461 ltor_fields = self._build_ltor_fields(tokens, labels)
463 # For padded sequences, ensure the embedding layer can map the token ID
464 tokens[tokens == self._pad_token_id] = 0
465 labels[labels == self._pad_token_id] = 0
467 if np.any(tokens < 0) or np.any(tokens >= len(self.config.tokenizer)):
468 raise ValueError("An input token is out of bounds of the tokenizer vocabulary")
470 raw_sample = {
471 "tokens": tokens,
472 "labels": labels,
473 }
474 raw_sample.update(ltor_fields)
476 return raw_sample
478 def _query_document_sample_shuffle_indices(
479 self,
480 index: int,
481 ) -> tuple[np.ndarray, np.ndarray]:
482 """Get the text (token ids) and document ids for a given index
484 Args:
485 index (int): The index into the dataset
487 Returns:
488 The text ids, document ids, and per-document token counts (before any padding).
489 """
490 shuffled_index = int(self.shuffle_index[index])
491 begin_document, begin_offset = self.sample_index[shuffled_index]
492 end_document, end_offset = self.sample_index[shuffled_index + 1]
494 document_ids = []
495 sample_parts = []
497 # Read one contiguous slice when the sample stays within a document.
498 if begin_document == end_document:
499 document_id = int(self.document_index[begin_document])
500 document_ids.append(document_id)
501 sample_parts.append(
502 self.dataset.get(
503 document_id,
504 offset=int(begin_offset),
505 length=int(end_offset - begin_offset + self.config.add_extra_token_to_sequence),
506 )
507 )
509 # Concatenate slices when the sample crosses document boundaries.
510 else:
511 for document_position in range(int(begin_document), int(end_document) + 1):
512 document_id = int(self.document_index[document_position])
513 document_ids.append(document_id)
514 offset = 0 if document_position > begin_document else int(begin_offset)
515 length = (
516 None
517 if document_position < end_document
518 else int(end_offset + self.config.add_extra_token_to_sequence)
519 )
520 sample_parts.append(self.dataset.get(document_id, offset=offset, length=length))
522 # Retained partial samples still need the configured fixed length.
523 length = sum(map(len, sample_parts))
524 document_lengths = [len(p) for p in sample_parts]
526 if length < (self.config.sequence_length + self.config.add_extra_token_to_sequence):
527 sample_parts.append(
528 [self._pad_token_id]
529 * (self.config.sequence_length + self.config.add_extra_token_to_sequence - length)
530 )
532 text = np.asarray(np.concatenate(sample_parts), dtype=np.int64)
534 return text, np.asarray(document_ids, dtype=np.int64), document_lengths
537class GPTFromMRDataset(_IndexedPretrainDataset):
538 """GPT Dataset for pre-packed samples converted from MindRecord.
540 Each entry in the underlying indexed Dataset represents a complete pre-packed
541 training sample. Samples are read directly without rebuilding document, sample,
542 or shuffle indices.
544 Args:
545 dataset: The indexed Dataset containing pre-packed samples converted from MindRecord.
547 dataset_path: The path to the underlying indexed Dataset.
549 indices: The indices of the pre-packed samples exposed by this Dataset.
551 num_samples: The number of pre-packed samples to draw from the indexed Dataset.
553 index_split: The Dataset split associated with ``indices``.
555 config: The GPT-specific Dataset configuration.
556 """
558 def _finalize(self) -> None:
559 """Validate the low-level Dataset and the pre-cut record lengths."""
560 if self.dataset is None or self.indices is None:
561 raise ValueError("GPTFromMRDataset requires a low-level Dataset and split indices")
563 if not isinstance(self.config, GPTDatasetConfig):
564 raise ValueError("GPTFromMRDataset requires GPTDatasetConfig")
566 if not self.config.skip_data_check:
567 expected_length = self.config.sequence_length + self.config.add_extra_token_to_sequence
568 if not np.all(self.dataset.sequence_lengths == expected_length):
569 raise ValueError(
570 "All pre-cut records must match sequence_length + add_extra_token_to_sequence; "
571 f"expected {expected_length}, got {self.dataset.sequence_lengths.tolist()}"
572 )
573 if self.num_samples > len(self.indices):
574 raise ValueError(
575 "Requested samples exceed the pre-cut records in this split: "
576 f"requested {self.num_samples}, available {len(self.indices)}"
577 )
579 @staticmethod
580 def numel_low_level_dataset(low_level_dataset: IndexedDataReader) -> int:
581 """Return the number of pre-cut records available for splitting."""
582 return int(low_level_dataset.sequence_lengths.shape[0])
584 @staticmethod
585 def is_multimodal() -> bool:
586 """Return whether records contain multimodal sequence modes."""
587 return False
589 @staticmethod
590 def is_split_by_sequence() -> bool:
591 """Return whether train/validation/test boundaries use sequence indices."""
592 return True
594 @staticmethod
595 def build_low_level_dataset(
596 dataset_path: str,
597 config: GPTDatasetConfig,
598 ) -> IndexedDataReader:
599 """Open a low-level reader for pre-cut records."""
600 return IndexedDataReader(
601 dataset_path,
602 mmap=config.mmap_bin_files,
603 reuse_index=config.reuse_idx,
604 )
606 def __len__(self) -> int:
607 """Return the requested number of pre-cut records in this split."""
608 return self.num_samples
610 def __getitem__(self, index: int) -> Mapping[str, Any]:
611 """Read and shift one pre-cut record without dynamic concatenation."""
612 sequence_id = int(self.indices[index])
613 text = np.asarray(self.dataset[sequence_id], dtype=np.int64)
614 if np.any(text == self._pad_token_id):
615 raise ValueError("Pre-cut indexed Dataset records must be packed without padding")
617 if self.config.add_extra_token_to_sequence:
618 tokens = text[:-1].copy()
619 labels = text[1:].copy()
620 else:
621 tokens = text.copy()
622 labels = np.roll(text, shift=-1)
623 labels[-1] = self._pad_token_id
625 if np.any(tokens < 0) or np.any(tokens >= self.config.tokenizer.vocab_size):
626 raise ValueError("An input token is out of bounds of the tokenizer vocabulary")
628 ltor_fields = self._build_ltor_fields(tokens, labels)
630 tokens[tokens == self._pad_token_id] = 0
631 labels[labels == self._pad_token_id] = 0
633 raw_sample = {
634 "tokens": tokens,
635 "labels": labels,
636 }
637 raw_sample.update(ltor_fields)
639 return raw_sample
642def _get_ltor_masks_and_position_ids(
643 tokens: np.ndarray,
644 config: GPTDatasetConfig,
645) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
646 """Build causal attention, loss, and position arrays for one token sequence.
648 Args:
649 tokens: Token IDs for one fixed-length Indexed sample.
650 config: Indexed Dataset runtime-field configuration.
652 Returns:
653 Attention mask, loss mask, and position IDs.
654 """
655 seq_length = tokens.size
656 eod_token = config.tokenizer.eod
657 attention_mask = np.tril(np.ones((seq_length, seq_length)))[None, :, :]
659 # Loss mask.
660 loss_mask = np.ones(seq_length, dtype=np.float32)
661 if config.eod_mask_loss:
662 loss_mask[tokens == eod_token] = 0.0
664 # Position ids.
665 position_ids = np.arange(seq_length, dtype=np.int64)
666 # We need to clone as the ids will be modified based on batch index.
667 if config.reset_position_ids:
668 position_ids = position_ids.copy()
670 if config.reset_position_ids or config.reset_attention_mask:
671 # Find indices where EOD token is.
672 eod_indices = position_ids[tokens == eod_token]
673 # Detach indices from positions if going to modify positions.
674 if config.reset_position_ids:
675 eod_indices = eod_indices.copy()
677 # Loop through EOD indices:
678 prev_index = 0
679 for eod_index in eod_indices:
680 next_index = int(eod_index) + 1
681 # Mask attention loss.
682 if config.reset_attention_mask:
683 attention_mask[0, next_index:, :next_index] = 0
684 # Reset positions.
685 if config.reset_position_ids:
686 position_ids[next_index:] -= next_index - prev_index
687 prev_index = next_index
689 # Convert attention mask to binary:
690 attention_mask = attention_mask < 0.5
692 return attention_mask, loss_mask, position_ids