| """I-JEPA multi-block masking for 1D token sequences (Weimann-style).""" |
| from __future__ import annotations |
|
|
| import torch |
|
|
|
|
| def multi_block_mask_1d( |
| n_tokens: int, |
| n_targets: int = 4, |
| target_size_range: tuple[int, int] = (4, 8), |
| mask_ratio: float = 0.5, |
| generator: torch.Generator | None = None, |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| """Return (context_idx, target_idx) for one sequence. |
| |
| Chooses `n_targets` contiguous blocks as targets (no overlap), then the |
| complement of their union is the context. mask_ratio caps total target |
| fraction. |
| """ |
| target_mask = torch.zeros(n_tokens, dtype=torch.bool) |
| max_cover = int(mask_ratio * n_tokens) |
| covered = 0 |
| attempts = 0 |
| while covered < max_cover and attempts < 64: |
| attempts += 1 |
| lo, hi = target_size_range |
| size = int(torch.randint(lo, hi + 1, (1,), generator=generator).item()) |
| size = min(size, max_cover - covered) |
| if size <= 0: |
| break |
| start = int(torch.randint(0, max(1, n_tokens - size + 1), (1,), generator=generator).item()) |
| if target_mask[start : start + size].any(): |
| continue |
| target_mask[start : start + size] = True |
| covered += size |
|
|
| target_idx = torch.nonzero(target_mask, as_tuple=False).squeeze(-1) |
| context_idx = torch.nonzero(~target_mask, as_tuple=False).squeeze(-1) |
| return context_idx, target_idx |
|
|
