| """Deterministic padded sequence samples for transformer models.""" |
|
|
| from __future__ import annotations |
|
|
| from collections.abc import Iterator, Sequence |
| from dataclasses import dataclass |
| from typing import Any |
|
|
| import torch |
|
|
| from gnn4colliders.data.events import PackedEvents |
| from gnn4colliders.data.metadata import BatchMetadata, EventMetadata |
| from gnn4colliders.data.sample import EventSample |
| from gnn4colliders.features import build_node_features |
|
|
|
|
| @dataclass(frozen=True) |
| class SequenceSample: |
| """One event represented as an ordered sequence of object tokens.""" |
|
|
| tokens: torch.Tensor |
| label: torch.Tensor |
| global_features: torch.Tensor | None |
| metadata: EventMetadata |
|
|
|
|
| @dataclass(frozen=True) |
| class SequenceBatch: |
| """A padded batch with a boolean mask of valid tokens.""" |
|
|
| tokens: torch.Tensor |
| token_mask: torch.Tensor |
| labels: torch.Tensor |
| global_features: torch.Tensor | None |
| metadata: BatchMetadata |
|
|
| def to( |
| self, device: torch.device | str, *, non_blocking: bool = False |
| ) -> "SequenceBatch": |
| return self._map(lambda value: value.to(device, non_blocking=non_blocking)) |
|
|
| def pin_memory(self) -> "SequenceBatch": |
| return self._map(lambda value: value.pin_memory()) |
|
|
| def _map(self, move: Any) -> "SequenceBatch": |
| metadata = BatchMetadata( |
| fold=move(self.metadata.fold), |
| weight=move(self.metadata.weight), |
| sample_id=self.metadata.sample_id, |
| extra=self.metadata.extra, |
| ) |
| return SequenceBatch( |
| tokens=move(self.tokens), |
| token_mask=move(self.token_mask), |
| labels=move(self.labels), |
| global_features=( |
| None if self.global_features is None else move(self.global_features) |
| ), |
| metadata=metadata, |
| ) |
|
|
|
|
| def sequence_sample_from_event( |
| event: EventSample, |
| feature_branches: Sequence[Any], |
| object_types: Sequence[str], |
| scales: Sequence[float] | torch.Tensor, |
| ) -> SequenceSample: |
| """Build tokens from shared collider features without graph topology.""" |
| tokens, _ = build_node_features( |
| event.objects, feature_branches, object_types, scales |
| ) |
| globals_ = torch.as_tensor(event.global_features, dtype=torch.float32) |
| return SequenceSample( |
| tokens=tokens, |
| label=torch.as_tensor(event.label), |
| global_features=globals_ if globals_.numel() else None, |
| metadata=event.event_metadata, |
| ) |
|
|
|
|
| def collate_sequence_batch( |
| events: PackedEvents, indices: torch.Tensor |
| ) -> SequenceBatch: |
| """Pad the object rows of ``indices`` into a token batch, in order.""" |
| if indices.numel() == 0: |
| raise ValueError("cannot batch an empty sequence") |
| rows, counts = events.gather_objects(indices) |
| max_length = int(counts.max()) if counts.numel() else 0 |
| batch_index = torch.repeat_interleave( |
| torch.arange(counts.numel()), counts, output_size=rows.shape[0] |
| ) |
| position = ( |
| torch.arange(rows.shape[0]) - (torch.cumsum(counts, 0) - counts)[batch_index] |
| ) |
| tokens = rows.new_zeros((counts.numel(), max_length, rows.shape[1])) |
| tokens[batch_index, position] = rows |
| token_mask = torch.zeros((counts.numel(), max_length), dtype=torch.bool) |
| token_mask[batch_index, position] = True |
| return SequenceBatch( |
| tokens=tokens, |
| token_mask=token_mask, |
| labels=events.labels[indices], |
| global_features=( |
| None if events.global_features is None else events.global_features[indices] |
| ), |
| metadata=events.batch_metadata(indices), |
| ) |
|
|
|
|
| def batch_sequence_samples(samples: Sequence[SequenceSample]) -> SequenceBatch: |
| """Pad a sequence of variable-length events in deterministic order.""" |
| if not samples: |
| raise ValueError("cannot batch an empty sequence") |
| feature_widths = {int(sample.tokens.shape[1]) for sample in samples} |
| if any(sample.tokens.ndim != 2 for sample in samples): |
| raise ValueError("sequence tokens must have shape [objects, features]") |
| if len(feature_widths) != 1: |
| raise ValueError("all sequence samples must have the same feature width") |
| max_length = max(int(sample.tokens.shape[0]) for sample in samples) |
| feature_width = next(iter(feature_widths)) |
| tokens = samples[0].tokens.new_zeros((len(samples), max_length, feature_width)) |
| token_mask = torch.zeros( |
| (len(samples), max_length), dtype=torch.bool, device=tokens.device |
| ) |
| for index, sample in enumerate(samples): |
| length = int(sample.tokens.shape[0]) |
| tokens[index, :length] = sample.tokens |
| token_mask[index, :length] = True |
| labels = torch.stack([torch.as_tensor(sample.label) for sample in samples]) |
| globals_ = [sample.global_features for sample in samples] |
| if any(value is None for value in globals_): |
| if not all(value is None for value in globals_): |
| raise ValueError("global_features must be present for every sample or none") |
| batched_globals = None |
| else: |
| batched_globals = torch.stack( |
| [value for value in globals_ if value is not None] |
| ) |
| return SequenceBatch( |
| tokens=tokens, |
| token_mask=token_mask, |
| labels=labels, |
| global_features=batched_globals, |
| metadata=BatchMetadata.from_events([sample.metadata for sample in samples]), |
| ) |
|
|
|
|
| class SequenceDataLoader: |
| """Small deterministic loader for variable-length sequence samples.""" |
|
|
| def __init__( |
| self, |
| dataset: Sequence[SequenceSample], |
| batch_size: int, |
| *, |
| shuffle: bool = False, |
| seed: int = 12345, |
| ) -> None: |
| if batch_size < 1: |
| raise ValueError("batch_size must be positive") |
| self.dataset = dataset |
| self.batch_size = batch_size |
| self.shuffle = shuffle |
| self.seed = seed |
| self.epoch = 0 |
|
|
| def set_epoch(self, epoch: int) -> None: |
| self.epoch = epoch |
|
|
| def __len__(self) -> int: |
| return (len(self.dataset) + self.batch_size - 1) // self.batch_size |
|
|
| def __iter__(self) -> Iterator[SequenceBatch]: |
| indices = torch.arange(len(self.dataset), dtype=torch.long) |
| if self.shuffle: |
| generator = torch.Generator().manual_seed(self.seed + self.epoch) |
| indices = indices[torch.randperm(len(indices), generator=generator)] |
| for start in range(0, len(indices), self.batch_size): |
| selected = [ |
| self.dataset[int(index)] |
| for index in indices[start : start + self.batch_size] |
| ] |
| yield batch_sequence_samples(selected) |
|
|
