File size: 2,445 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 68 69 70 71 72 73 | # 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
import torch.nn.functional as functional
def supervised_token_count(
labels: torch.Tensor | None,
*,
label_offset: int = 0,
ignore_index: int = -100,
) -> torch.Tensor:
if labels is None or not torch.is_tensor(labels):
return torch.zeros((), dtype=torch.int64)
if labels.ndim != 2:
raise ValueError(
f"labels must be 2D [B,S] to count supervised tokens (got shape={tuple(labels.shape)})"
)
target_start = int(label_offset) + 1
if int(labels.size(1)) <= target_start:
return torch.zeros((), device=labels.device, dtype=torch.int64)
shifted = labels[:, target_start:]
return (shifted != int(ignore_index)).sum(dtype=torch.int64)
def shifted_loss_sum_and_count(
logits: torch.Tensor,
labels: torch.Tensor,
*,
label_offset: int = 0,
ignore_index: int = -100,
) -> tuple[torch.Tensor, torch.Tensor]:
if int(logits.size(1)) <= 1 or int(labels.size(1)) <= int(label_offset) + 1:
zero = logits.new_zeros(())
return zero, zero
shift_logits = logits[:, :-1, :].contiguous()
target_start = int(label_offset) + 1
target_end = target_start + int(shift_logits.size(1))
if int(labels.size(1)) < target_end:
shift_logits = shift_logits[:, : max(int(labels.size(1)) - target_start, 0), :]
target_end = target_start + int(shift_logits.size(1))
if int(shift_logits.size(1)) <= 0:
zero = logits.new_zeros(())
return zero, zero
shift_labels = labels[:, target_start:target_end].contiguous()
flat_labels = shift_labels.reshape(-1).to(dtype=torch.long)
flat_logits = shift_logits.reshape(-1, int(shift_logits.size(-1)))
loss_sum = functional.cross_entropy(
flat_logits,
flat_labels,
ignore_index=int(ignore_index),
reduction="sum",
)
count = (flat_labels != int(ignore_index)).sum().to(dtype=loss_sum.dtype)
return loss_sum, count
def mean_loss_from_sum_and_count(
*,
loss_sum: torch.Tensor,
count: torch.Tensor,
reference: torch.Tensor,
) -> torch.Tensor:
return torch.where(
count > 0,
loss_sum / count.clamp_min(1.0),
reference.new_zeros(()),
)
|