| """Peptide length constants and class-index mapping.""" |
|
|
| from __future__ import annotations |
|
|
| import torch |
|
|
| MIN_PEPTIDE_LENGTH = 1 |
| MAX_PEPTIDE_LENGTH = 30 |
| NUM_LENGTH_CLASSES = MAX_PEPTIDE_LENGTH - MIN_PEPTIDE_LENGTH + 1 |
|
|
|
|
| def validate_peptide_length(length: int) -> None: |
| if not MIN_PEPTIDE_LENGTH <= length <= MAX_PEPTIDE_LENGTH: |
| raise ValueError( |
| f"Peptide length {length} is outside supported range " |
| f"[{MIN_PEPTIDE_LENGTH}, {MAX_PEPTIDE_LENGTH}]." |
| ) |
|
|
|
|
| def length_to_class(length: torch.Tensor) -> torch.Tensor: |
| """Map peptide length L ∈ [1, 30] to class index L - 1 ∈ [0, 29].""" |
| return length - MIN_PEPTIDE_LENGTH |
|
|
|
|
| def class_to_length(class_id: torch.Tensor) -> torch.Tensor: |
| """Map class index ∈ [0, 29] back to peptide length ∈ [1, 30].""" |
| return class_id + MIN_PEPTIDE_LENGTH |
|
|
|
|
| def clamp_length(length: torch.Tensor) -> torch.Tensor: |
| return length.clamp(min=MIN_PEPTIDE_LENGTH, max=MAX_PEPTIDE_LENGTH) |
|
|
|
|
| def length_to_active_mask(length: torch.Tensor, seq_len: int) -> torch.Tensor: |
| """True for positions 0 .. length-1 (peptide residues), False for padding tail.""" |
| positions = torch.arange(seq_len, device=length.device) |
| return positions.unsqueeze(0) < length.unsqueeze(1) |
|
|
|
|
| def apply_length_padding( |
| token_ids: torch.Tensor, |
| length: torch.Tensor, |
| pad_id: int, |
| ) -> torch.Tensor: |
| """Force `<pad>` at all positions j >= length (structural padding, not sampled).""" |
| active_mask = length_to_active_mask(length, token_ids.shape[1]) |
| return token_ids.masked_fill(~active_mask, pad_id) |
|
|
|
|
