File size: 2,302 Bytes
d53adc9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 | # Generated by ml.integrations.export.runtime_packager.write_remote_code_bundle.
# Exported for HuggingFace trust_remote_code loading.
# This file is intentionally self-contained.
from __future__ import annotations
import torch
def _as_bool_mask(attention_mask: torch.Tensor) -> torch.Tensor:
return (
attention_mask if attention_mask.dtype == torch.bool else (attention_mask != 0)
)
def is_all_ones_mask(attention_mask: torch.Tensor) -> bool:
mask = _as_bool_mask(attention_mask)
return bool(mask.all().item())
def validate_right_padding_mask(
attention_mask: torch.Tensor,
*,
input_ids: torch.Tensor,
) -> None:
if attention_mask.dim() != 2:
raise ValueError("attention_mask must be [B,T]")
if attention_mask.shape != input_ids.shape:
raise ValueError("attention_mask shape must match input_ids")
mask = _as_bool_mask(attention_mask)
valid_rows = mask.any(dim=1)
resumes_after_padding = (~mask[:, :-1] & mask[:, 1:]).any(dim=1)
if not bool(valid_rows.all().item()):
raise ValueError("attention_mask row has no valid tokens")
if bool(resumes_after_padding.any().item()):
raise ValueError("loss computation requires right-padded attention masks")
def slice_valid_tokens(
input_ids: torch.Tensor,
attention_mask: torch.Tensor,
) -> list[tuple[int, int, torch.Tensor]]:
if attention_mask.dim() != 2:
raise ValueError("attention_mask must be [B,T]")
mask = _as_bool_mask(attention_mask)
if mask.shape != input_ids.shape:
raise ValueError("attention_mask shape must match input_ids")
rows: list[tuple[int, int, torch.Tensor]] = []
for batch_index in range(int(input_ids.size(0))):
idx = mask[batch_index].nonzero(as_tuple=False).squeeze(-1)
if int(idx.numel()) == 0:
raise ValueError("attention_mask row has no valid tokens")
start = int(idx[0].item())
end = int(idx[-1].item()) + 1
if int(idx.numel()) != (end - start):
raise ValueError("Sophia only supports contiguous padding masks")
rows.append((start, end, input_ids[batch_index : batch_index + 1, start:end]))
return rows
__all__ = [
"is_all_ones_mask",
"slice_valid_tokens",
"validate_right_padding_mask",
]
|