File size: 1,323 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 | # 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
from .input_mask import validate_right_padding_mask
from .loss_stats import mean_loss_from_sum_and_count
from .decoder_types import DecoderConfig
from .decoder_types import DecoderCoreModel
from .decoder_loss import chunked_loss_stats_from_hidden
def forward_loss(
*,
runtime_model: DecoderCoreModel,
config: DecoderConfig,
training: bool,
input_ids: torch.Tensor,
attention_mask: torch.Tensor | None,
labels: torch.Tensor,
output_weight: torch.Tensor,
) -> torch.Tensor:
del training
ref = output_weight
if attention_mask is not None:
validate_right_padding_mask(attention_mask, input_ids=input_ids)
hidden, _ = runtime_model._forward_hidden(input_ids, start_pos=0)
base_sum, base_count = chunked_loss_stats_from_hidden(
hidden,
labels,
label_offset=0,
norm=runtime_model.norm,
output=runtime_model.output,
config=config,
)
return mean_loss_from_sum_and_count(
loss_sum=base_sum,
count=base_count,
reference=ref,
)
__all__ = ["forward_loss"]
|