39 lines
1.7 KiB
Python
39 lines
1.7 KiB
Python
import torch
|
|
|
|
|
|
def validate_prompt_sequence_lengths(
|
|
attention_mask: torch.Tensor,
|
|
*,
|
|
max_sequence_length: int,
|
|
supported_max_sequence_length: int,
|
|
prompt_name: str = "prompt",
|
|
length_offset: int = 0,
|
|
baseline_attention_mask: torch.Tensor | None = None,
|
|
error_context: str,
|
|
) -> None:
|
|
sequence_lengths = attention_mask.sum(dim=1)
|
|
if baseline_attention_mask is not None:
|
|
# Some callers need to validate only the user-controlled portion of a
|
|
# templated prompt. In those cases we subtract the fully-tokenized
|
|
# template baseline instead of only removing a fixed prefix length,
|
|
# because the template may also contribute a suffix or image markers.
|
|
baseline_lengths = baseline_attention_mask.sum(dim=1)
|
|
if baseline_lengths.shape[0] == 1 and sequence_lengths.shape[0] > 1:
|
|
baseline_lengths = baseline_lengths.expand(sequence_lengths.shape[0])
|
|
sequence_lengths = sequence_lengths - baseline_lengths
|
|
if length_offset:
|
|
sequence_lengths = sequence_lengths - length_offset
|
|
sequence_lengths = torch.clamp(sequence_lengths, min=0)
|
|
too_long = torch.nonzero(sequence_lengths > max_sequence_length, as_tuple=False)
|
|
if too_long.numel() == 0:
|
|
return
|
|
|
|
batch_idx = int(too_long[0].item())
|
|
actual_length = int(sequence_lengths[batch_idx].item())
|
|
prompt_ref = f"`{prompt_name}` at batch index {batch_idx}" if attention_mask.shape[0] > 1 else f"`{prompt_name}`"
|
|
raise ValueError(
|
|
f"{prompt_ref} is too long {error_context}: got {actual_length} tokens, but "
|
|
f"`max_sequence_length` is {max_sequence_length}. Shorten the prompt or increase "
|
|
f"`max_sequence_length` up to {supported_max_sequence_length}."
|
|
)
|