Randomly masks contiguous sections of an unbatched sequence, ensuring separation between chunks is at least chunk_width.
(self, seq: torch.Tensor)
| 153 | return torch.cat(sections, 1) |
| 154 | |
| 155 | def _mask_seq(self, seq: torch.Tensor) -> torch.Tensor: |
| 156 | """Randomly masks contiguous sections of an unbatched sequence, |
| 157 | ensuring separation between chunks is at least chunk_width.""" |
| 158 | len_ = seq.shape[0] |
| 159 | num_chunks = self.hparams.mask_num_chunks |
| 160 | chunk_width = self.hparams.mask_chunk_width |
| 161 | |
| 162 | # Ensure there's enough space for the chunks and separations |
| 163 | total_width_needed = num_chunks * chunk_width + (num_chunks - 1) * chunk_width |
| 164 | if total_width_needed > len_: |
| 165 | raise ValueError("Sequence is too short to mask") |
| 166 | |
| 167 | masked_seq = seq.clone() |
| 168 | |
| 169 | for i in range(num_chunks): |
| 170 | start = (i * len_) // num_chunks |
| 171 | loc = torch.randint(0, len_ // num_chunks - chunk_width, (1,)).item() |
| 172 | masked_seq[loc + start : loc + start + chunk_width] = 0 |
| 173 | |
| 174 | return masked_seq |