| """Dynamic padding + next-token label construction.""" | |
| from __future__ import annotations | |
| import numpy as np | |
| import torch | |
| from ..tokenizers.multimodal import Sample | |
| def collate(samples: list[Sample], pad_id: int, multiple: int = 8, *, target_length: int | None = None, | |
| batch_size: int | None = None) -> tuple[torch.Tensor, torch.Tensor]: | |
| """Right-pad to the batch max (rounded up to `multiple`) and build shifted labels. | |
| labels[j] = tokens[j+1] where j+1 >= target_start, else -100 (prompt and padding are not | |
| trained on). Right padding is safe under causal attention: real tokens never attend to it. | |
| """ | |
| L = max(len(s.tokens) for s in samples) - 1 | |
| if target_length is not None: | |
| if target_length < L: | |
| raise ValueError("fixed padding cannot truncate a sample") | |
| L = target_length | |
| else: | |
| L = -(-L // multiple) * multiple | |
| B = len(samples) if batch_size is None else batch_size | |
| if B < len(samples): | |
| raise ValueError("fixed batch size cannot discard samples") | |
| inp = np.full((B, L), pad_id, dtype=np.int64) | |
| lab = np.full((B, L), -100, dtype=np.int64) | |
| for i, s in enumerate(samples): | |
| t = s.tokens.astype(np.int64) | |
| n = len(t) - 1 | |
| inp[i, :n] = t[:-1] | |
| lab[i, :n] = t[1:] | |
| lab[i, : max(0, s.target_start - 1)] = -100 | |
| return torch.from_numpy(inp), torch.from_numpy(lab) | |