| import random |
| import torch |
| from data.lengths import MAX_PEPTIDE_LENGTH, MIN_PEPTIDE_LENGTH |
|
|
|
|
| def length_noiser( |
| batch_lengths: torch.Tensor, |
| noising_prob: float = 0.1, |
| max_: int = MAX_PEPTIDE_LENGTH, |
| min_: int = MIN_PEPTIDE_LENGTH, |
| ) -> tuple[torch.Tensor, bool]: |
| """ |
| Perturb batch peptide lengths by +/-1 with probability `noising_prob`. |
| |
| Args: |
| batch_lengths: Tensor of peptide lengths (B,). |
| noising_prob: Probability of perturbing lengths. |
| max_: Maximum allowed peptide length. |
| min_: Minimum allowed peptide length. |
| |
| Returns: |
| (noised_lengths, is_noised) |
| """ |
| r = random.random() |
| if r <= noising_prob / 2.0: |
| noised = (batch_lengths - 1).clamp(min=min_, max=max_) |
| is_noised = bool((noised != batch_lengths).any().item()) |
| return noised, is_noised |
|
|
| elif r <= noising_prob: |
| noised = (batch_lengths + 1).clamp(min=min_, max=max_) |
| is_noised = bool((noised != batch_lengths).any().item()) |
| return noised, is_noised |
|
|
| return batch_lengths, False |
|
|
