| """Deterministic prefix crops over the unchanged sentence window token store.""" |
|
|
| import math |
|
|
| import numpy as np |
|
|
| from vimeml.training.data import SentenceWindowDataset |
|
|
| MASK64 = (1 << 64) - 1 |
| GOLDEN64 = 0x9E3779B97F4A7C15 |
| EPOCH64 = 0xD1B54A32D192ED03 |
|
|
|
|
| def mix64(value): |
| value = ((value ^ (value >> 30)) * 0xBF58476D1CE4E5B9) & MASK64 |
| value = ((value ^ (value >> 27)) * 0x94D049BB133111EB) & MASK64 |
| return value ^ (value >> 31) |
|
|
|
|
| def mix64_array(values): |
| values = (values ^ (values >> np.uint64(30))) * np.uint64(0xBF58476D1CE4E5B9) |
| values = (values ^ (values >> np.uint64(27))) * np.uint64(0x94D049BB133111EB) |
| return values ^ (values >> np.uint64(31)) |
|
|
|
|
| class PrefixCropWindowDataset(SentenceWindowDataset): |
| |
| epoch_indexed = True |
|
|
| def __init__( |
| self, |
| token_dir, |
| index_dir, |
| split="train", |
| *, |
| seed=42, |
| probability=0.30, |
| min_remaining_tokens=8, |
| ): |
| if split != "train": |
| raise ValueError("Prefix crops are only defined for the V2 training split.") |
| if not math.isfinite(probability) or not 0 <= probability <= 1: |
| raise ValueError("Crop probability must be in [0, 1].") |
| if type(min_remaining_tokens) is not int or min_remaining_tokens < 1: |
| raise ValueError("Minimum remaining tokens must be a positive integer.") |
| if type(seed) is not int or not 0 <= seed <= MASK64: |
| raise ValueError("Crop seed must be a uint64 integer.") |
| super().__init__(token_dir, index_dir, split) |
| self.seed = seed |
| self.probability = probability |
| self.threshold = int(probability * (1 << 53)) |
| self.min_remaining_tokens = min_remaining_tokens |
| self.epoch = 0 |
|
|
| def set_epoch(self, epoch): |
| if type(epoch) is not int or epoch < 0: |
| raise ValueError("Crop epoch must be a nonnegative integer.") |
| self.epoch = epoch |
|
|
| def crop_offset(self, index, length, original_start, epoch): |
| maximum = length - self.min_remaining_tokens |
| if original_start != 0 or maximum < 1 or self.threshold == 0: |
| return 0 |
| bits = mix64((self.seed + epoch * EPOCH64 + index + GOLDEN64) & MASK64) |
| if (bits >> 11) >= self.threshold: |
| return 0 |
| return 1 + mix64((bits + GOLDEN64) & MASK64) % maximum |
|
|
| def __getitem__(self, key): |
| epoch, index = key if isinstance(key, tuple) else (self.epoch, key) |
| index = int(index) |
| if index < 0: |
| index += len(self) |
| if type(epoch) is not int or epoch < 0: |
| raise ValueError("Crop epoch must be a nonnegative integer.") |
| sample = super().__getitem__(index) |
| original_start = sample["window_start"] |
| length = len(sample["input_ids"]) |
| offset = self.crop_offset(index, length, original_start, epoch) |
| return { |
| **sample, |
| "input_ids": sample["input_ids"][offset:], |
| "labels": sample["labels"][offset:], |
| "window_start": original_start + offset, |
| "original_window_start": original_start, |
| "crop_offset": offset, |
| "uncropped_length": length, |
| "sample_index": index, |
| "sample_epoch": epoch, |
| } |
|
|
| def window_lengths(self, indices): |
| indices = np.asarray(indices, dtype=np.int64) |
| lengths = super().window_lengths(indices) |
| if not self.threshold: |
| return lengths |
| position = np.searchsorted(self.end, indices, side="right") |
| first_window = np.ones(len(indices), dtype=np.bool_) |
| candidates = np.flatnonzero(position < len(self.end)) |
| first_window[candidates] = indices[candidates] <= self.first[position[candidates]] |
| maximum = lengths - self.min_remaining_tokens |
| base = (self.seed + self.epoch * EPOCH64 + GOLDEN64) & MASK64 |
| bits = mix64_array(indices.astype(np.uint64) + np.uint64(base)) |
| apply = first_window & (maximum > 0) & ((bits >> np.uint64(11)) < np.uint64(self.threshold)) |
| choices = mix64_array(bits + np.uint64(GOLDEN64)) % np.maximum(maximum, 1).astype( |
| np.uint64 |
| ) + np.uint64(1) |
| return lengths - np.where(apply, choices, 0).astype(np.int64) |
|
|
