import torch import torch.nn as nn import torch.nn.functional as F from .layers import ( AdaNLayers, AdaLNZero, MzArrayEncoder, SinusoidalEmbedding, SwiGLU, SwiGLUFFN, modulate, ) from config.defaults import MODEL_CONFIG from data.lengths import MAX_PEPTIDE_LENGTH, MIN_PEPTIDE_LENGTH, NUM_LENGTH_CLASSES class SpectrumEncoder(nn.Module): def __init__( self, model_dim: int = 512, num_layers: int = 6, nhead: int = 8, dim_feedforward: int = 1536, dropout: float = 0.1, ): super().__init__() self.mz_encoder = MzArrayEncoder(emb_dim=model_dim // 2) self.comp_mz_encoder = MzArrayEncoder(emb_dim=model_dim // 2) self.intensity_proj = nn.Linear(1, model_dim) self.peak_proj = nn.Linear(model_dim * 2, model_dim) self.peak_norm = nn.LayerNorm(model_dim) encoder_layer = nn.TransformerEncoderLayer( d_model=model_dim, nhead=nhead, dim_feedforward=dim_feedforward, dropout=dropout, batch_first=True, norm_first=True, # Pre-LayerNorm avoids vanishing gradients in deep layers ) self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers) self.final_norm = nn.LayerNorm(model_dim) self.cls_token = nn.Parameter(torch.randn(model_dim)) def forward(self, mz1, mz2, intensity, spectrum_mask): """ Args: mz1: mass array (B, S) mz2: complementary mass array (B, S) intensity: peak intensities (B, S) spectrum_mask: padding mask where True marks padded positions (B, S) Returns: cls_embedding (B, D), peak_embeddings (B, S, D), peak_mask (B, S) """ batch_size = mz1.shape[0] device = mz1.device cls_token = self.cls_token.expand(batch_size, 1, -1) cls_mask = torch.zeros(batch_size, 1, dtype=torch.bool, device=device) full_mask = torch.cat((cls_mask, spectrum_mask), dim=1) # Fuse mz1, complementary mz2, and intensity per peak into a single D-dim token mz1_emb = self.mz_encoder(mz1) mz2_emb = self.comp_mz_encoder(mz2) mz_combined = torch.cat([mz1_emb, mz2_emb], dim=-1) int_emb = self.intensity_proj(intensity.unsqueeze(-1)) peak_tokens = self.peak_norm(self.peak_proj(torch.cat([mz_combined, int_emb], dim=-1))) x = torch.cat((cls_token, peak_tokens), dim=1) x = self.transformer_encoder(x, src_key_padding_mask=full_mask) x = self.final_norm(x) return x[:, 0, :], x[:, 1:, :], spectrum_mask class DecoderBlock(nn.Module): def __init__(self, emb_dim=512, num_heads=8, mlp_hidden_dim=1536, dropout=0.1): super().__init__() # 9 modulation parameters with zero-initialized gates (AdaLN-Zero): # (shift_msa, scale_msa, gate_msa, shift_cross, scale_cross, gate_cross, shift_mlp, scale_mlp, gate_mlp) self.ada_ln = AdaLNZero(emb_dim=emb_dim, num_params=9) self.self_attention = nn.MultiheadAttention( emb_dim, num_heads=num_heads, dropout=dropout, batch_first=True ) self.cross_attention = nn.MultiheadAttention( emb_dim, num_heads=num_heads, dropout=dropout, batch_first=True ) self.norm1 = nn.LayerNorm(emb_dim, elementwise_affine=False) self.norm2 = nn.LayerNorm(emb_dim, elementwise_affine=False) self.norm3 = nn.LayerNorm(emb_dim, elementwise_affine=False) self.mlp = SwiGLUFFN(d_model=emb_dim, d_ff=mlp_hidden_dim, dropout=dropout) def forward(self, conditioner, x, y, full_mask=None, sequence_padding_mask=None): ( shift_msa, scale_msa, gate_msa, shift_cross, scale_cross, gate_cross, shift_mlp, scale_mlp, gate_mlp, ) = self.ada_ln(conditioner) # 1. Self-Attention with AdaLN modulation and residual gate norm_x1 = modulate(self.norm1(x), shift_msa, scale_msa) attn_out, _ = self.self_attention( query=norm_x1, key=norm_x1, value=norm_x1, key_padding_mask=sequence_padding_mask, ) x = x + gate_msa.unsqueeze(1) * attn_out # 2. Cross-Attention with AdaLN modulation and residual gate norm_x2 = modulate(self.norm2(x), shift_cross, scale_cross) cross_mask = full_mask if cross_mask is not None and cross_mask.all(dim=-1).any(): cross_mask = cross_mask.clone() cross_mask[cross_mask.all(dim=-1), 0] = False cross_out, _ = self.cross_attention( query=norm_x2, key=y, value=y, key_padding_mask=cross_mask, ) x = x + gate_cross.unsqueeze(1) * cross_out # 3. Feed-Forward SwiGLU with AdaLN modulation and residual gate norm_x3 = modulate(self.norm3(x), shift_mlp, scale_mlp) x = x + gate_mlp.unsqueeze(1) * self.mlp(norm_x3) return x class DFMPeptideDecoder(nn.Module): def __init__( self, spec_dim: int = 512, emb_dim: int = 512, mlp_hidden_dim: int = 1536, vocab_size=22, mask_token_id=None, pad_token_id=None, max_charge=6, min_length: int = MIN_PEPTIDE_LENGTH, max_length: int = MAX_PEPTIDE_LENGTH, n_decoder_blocks: int = 6, num_heads: int = 8, dropout: float = 0.1, ): super().__init__() if mask_token_id is None: mask_token_id = vocab_size - 1 if pad_token_id is None: pad_token_id = vocab_size - 2 output_token_ids = [ token_id for token_id in range(vocab_size) if token_id not in {mask_token_id, pad_token_id} ] if output_token_ids != list(range(vocab_size - 2)): raise ValueError( "Decoder output logits require and to be the last two " f"vocabulary entries (pad={pad_token_id}, mask={mask_token_id}, " f"vocab_size={vocab_size})." ) self.vocab_size = vocab_size self.mask_token_id = mask_token_id self.pad_token_id = pad_token_id self.output_vocab_size = vocab_size - 2 self.sinusoidal_embedding = SinusoidalEmbedding(emb_dim) # Input embedding uses the full vocabulary (noisy x_t may contain ). self.peptide_embedding = nn.Embedding(num_embeddings=vocab_size, embedding_dim=emb_dim) # Charges are 1-indexed in the dataset; reserve index 0 for padding/unused. self.charge_embedding = nn.Embedding( num_embeddings=max_charge + 1, embedding_dim=emb_dim, ) self.min_length = min_length self.max_length = max_length # Lengths are in [min_length, max_length]; reserve indices 0..max_length for 1-based indexing. self.length_embedding = nn.Embedding( num_embeddings=max_length + 1, embedding_dim=emb_dim, ) self.decoder_blocks = nn.ModuleList( [ DecoderBlock( emb_dim=emb_dim, num_heads=num_heads, mlp_hidden_dim=mlp_hidden_dim, dropout=dropout, ) for _ in range(n_decoder_blocks) ] ) self.spectrum_proj = nn.Linear(spec_dim, emb_dim) self.conditioner_proj = nn.Linear(4 * emb_dim, emb_dim) self.norm = nn.LayerNorm(emb_dim, elementwise_affine=False) self.ada_final = nn.Sequential( nn.SiLU(), nn.Linear(emb_dim, 2 * emb_dim), ) nn.init.zeros_(self.ada_final[1].weight) nn.init.zeros_(self.ada_final[1].bias) # Output head: Linear projection from conditioned representation to vocabulary self.head = nn.Linear(emb_dim, self.output_vocab_size) @classmethod def from_vocabulary(cls, vocab: dict[str, int], **kwargs): """Build a decoder whose output head excludes the .""" mask_id = vocab.get("", vocab.get("")) if mask_id is None: raise KeyError("Vocabulary must contain either '' or ''.") pad_id = vocab[""] vocab_size = max(vocab.values()) + 1 return cls( vocab_size=vocab_size, mask_token_id=mask_id, pad_token_id=pad_id, **kwargs, ) def forward( self, time, precursor_m, precursor_c, spectrum_embeddings, peptide_seq, length, full_mask=None, sequence_padding_mask=None, ): seq_len = peptide_seq.shape[1] t = self.sinusoidal_embedding(time) peptide_emb = self.peptide_embedding(peptide_seq) positions = torch.arange(seq_len, device=peptide_seq.device, dtype=torch.float32) pos_emb = self.sinusoidal_embedding(positions).unsqueeze(0) charge_index = precursor_c.long().clamp(min=1, max=self.charge_embedding.num_embeddings - 1) c_emb = self.charge_embedding(charge_index) m_emb = self.sinusoidal_embedding(precursor_m) length_index = length.long().clamp(min=self.min_length, max=self.max_length) length_emb = self.length_embedding(length_index) if m_emb.dim() > 2: m_emb = m_emb.squeeze(1) if c_emb.dim() > 2: c_emb = c_emb.squeeze(1) if length_emb.dim() > 2: length_emb = length_emb.squeeze(1) if t.dim() > 2: t = t.squeeze(1) precursor_m_c_t_l = self.conditioner_proj( torch.cat((c_emb, m_emb, t, length_emb), dim=-1) ) x = peptide_emb + pos_emb y = self.spectrum_proj(spectrum_embeddings) for block in self.decoder_blocks: x = block( precursor_m_c_t_l, x, y, full_mask=full_mask, sequence_padding_mask=sequence_padding_mask, ) if sequence_padding_mask is not None: x = x.masked_fill(sequence_padding_mask.unsqueeze(-1), 0.0) shift_final, scale_final = self.ada_final(precursor_m_c_t_l).chunk(2, dim=-1) x = modulate(self.norm(x), shift_final, scale_final) # Output logits: amino acids only (no , no ). return self.head(x) class PeptideLengthClassifier(nn.Module): def __init__( self, n_classes: int = NUM_LENGTH_CLASSES, min_length: int = MIN_PEPTIDE_LENGTH, max_length: int = MAX_PEPTIDE_LENGTH, in_dim: int = 512, hidden_dim: int = 256, emb_dim: int = 128, max_charge: int = 10, dropout: float = 0.15, ): super().__init__() self.min_length = min_length self.max_length = max_length self.n_classes = n_classes self.mass_emb = SinusoidalEmbedding(emb_dim) self.charge_emb = nn.Embedding(max_charge + 1, emb_dim) self.mlp = nn.Sequential( nn.Linear(in_dim + 2 * emb_dim, hidden_dim), nn.SiLU(), nn.Dropout(dropout), nn.Linear(hidden_dim, hidden_dim * 2), nn.LayerNorm(hidden_dim * 2), nn.SiLU(), nn.Dropout(dropout), nn.Linear(hidden_dim * 2, n_classes), ) def forward(self, spectrum_embeddings, precursor_mass, precursor_charge): m = self.mass_emb(precursor_mass.reshape(-1)) charge_index = precursor_charge.long().clamp( min=0, max=self.charge_emb.num_embeddings - 1 ) c = self.charge_emb(charge_index) x = torch.cat([spectrum_embeddings.reshape(m.shape[0], -1), m, c], dim=-1) return self.mlp(x) # Backward-compatible aliases for older imports. DecoderBloc = DecoderBlock PeptideLenghtClassifier = PeptideLengthClassifier