Project: s41467-025-58866-4 | Function: forward (python)

system: You are an expert Python programmer implementing algorithms from scientific papers.
Complete the masked function based on the algorithm description.
Follow Python best practices and maintain consistent code style.
Return **only** the complete Python function code.

user: Complete this Python function:

**DESCRIPTION:** Main forward pass of the DIA-BERT transformer model. Processes mass spectrometry data through convolutional layers, dual transformer encoders, feature embeddings, and dense layers to produce peptide identification scores.

**KEY TERMS:** {
  "transformer_encoder_1": "transformer model first stage",
  "transformer_encoder_2": "transformer model second stage",
  "rsm": "mass spectrometry data",
  "dense_rep": "learned representation"
}

**CONTEXT CODE:**         rsm, rsm_max = FeatureEngineer.process_intensity_np(rsm)
        rsm = rsm.transpose(-2, -3)  # (B, 72, 8, 16) ==> (B, 8, 72, 16)

        feature_engineer = FeatureEngineer()
        frag_info = feature_engineer.process_frag_info(frag_info, max_intensity)

        # precursor_feat（B, 10）
        # 'sequence_length', 'precursor_mz', 'charge', 'precursor_irt', 'nr_peaks',
        # 'assay_rt_kept', 'delta_rt_kept', 'max_intensity', rt, instrument
        precursor_feat = np.column_stack((precursor_feat[:, :7], rsm_max))
        precursor_feat = feature_engineer.process_feat_np(precursor_feat)

        rt = feature_engineer.get_rt_s2i(rt)
        instrument = feature_engineer.get_rt_s2i(instrument)
        rt_tensor = torch.tensor(rt).repeat(frag_info.shape[0]).unsqueeze(dim=-1)
        instrument_tensor = torch.tensor(instrument).repeat(frag_info.shape[0]).unsqueeze(dim=-1)
        precursor_feat = torch.cat((precursor_feat, rt_tensor, instrument_tensor), dim=1)

        return rsm, frag_info, precursor_feat

    @staticmethod
    def feature_engineer(rsm, frag_info, precursor_feat, max_intensity=None, rt=26, instrument='Tripletof 6600'):
        rsm, rsm_max = FeatureEngineer.process_intensity(rsm)
        rsm = rsm.transpose(-2, -3)  # (B, 72, 8, 16) ==> (B, 8, 72, 16)

        feature_engineer = FeatureEngineer()
        frag_info = feature_engineer.process_frag_info(frag_info, max_intensity)

        # precursor_feat (B, 10)
        # 'sequence_length', 'precursor_mz', 'charge', 'precursor_irt', 'nr_peaks',
        # 'assay_rt_kept', 'delta_rt_kept', 'max_intensity', rt, instrument
        precursor_feat = torch.column_stack((precursor_feat[:, :7], rsm_max))
        precursor_feat = feature_engineer.process_feat(precursor_feat)

        rt = feature_engineer.get_rt_s2i(rt)
        instrument = feature_engineer.get_instrument_s2i(instrument)
        rt_tensor = torch.tensor(rt, dtype=torch.float32, device=precursor_feat.device).repeat(
            frag_info.shape[0]).unsqueeze(dim=-1)
        instrument_tensor = torch.tensor(instrument, dtype=torch.float32, device=precursor_feat.device).repeat(
            frag_info.shape[0]).unsqueeze(dim=-1)
        precursor_feat = torch.cat((precursor_feat, rt_tensor, instrument_tensor), dim=1)

        return rsm, frag_info, precursor_feat


class NLinearMemoryEfficient(nn.Module):
    def __init__(self, n: int, d_in: int, d_out: int) -> None:
        super().__init__()
        self.layers = nn.ModuleList([nn.Linear(d_in, d_out) for _ in range(n)])

    def forward(self, x):
        return torch.stack([l(x[:, i]) for i, l in enumerate(self.layers)], 1)


class NumEmbeddings(nn.Module):
    def __init__(
            self,
            n_features: int,
            d_embedding: int,
            embedding_arch: list,
            d_feature: int,
    ) -> None:
        super().__init__()
        assert embedding_arch
        assert set(embedding_arch) <= {
            'linear',
            'shared_linear',
            'relu',
            'layernorm',
            'batchnorm',
        }

        # NLinear_ =  NLinear
        layers: list[nn.Module] = []

        if embedding_arch[0] == 'linear':
            assert d_embedding is not None
            layers.append(
                NLinearMemoryEfficient(n_features, d_feature, d_embedding)
            )
        elif embedding_arch[0] == 'shared_linear':
            layers.append(
                nn.Linear(d_feature, d_embedding)
            )
        d_current = d_embedding

        for x in embedding_arch[1:]:
            layers.append(
                nn.ReLU()
                if x == 'relu'
                else NLinearMemoryEfficient(n_features, d_current, d_embedding)  # type: ignore[code]
                if x == 'linear'
                else nn.Linear(d_current, d_embedding)  # type: ignore[code]
                if x == 'shared_linear'
                else nn.LayerNorm([n_features, d_current])
                if x == 'layernorm'
                else nn.BatchNorm1d(n_features)
                if x == 'batchnorm'
                else nn.Identity()
            )
            if x in ['linear']:
                d_current = d_embedding
            assert not isinstance(layers[-1], nn.Identity)
        self.d_embedding = d_current
        self.layers = nn.Sequential(*layers)

    def forward(self, x):
        return self.layers(x)


class RMSNorm(torch.nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6):
        """
        Initialize the RMSNorm normalization layer.

        Args:
            dim (int): The dimension of the input tensor.
            eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6.

        Attributes:
            eps (float): A small value added to the denominator for numerical stability.
            weight (nn.Parameter): Learnable scaling parameter.

        """
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def _norm(self, x):
        """
        Apply the RMSNorm normalization to the input tensor.

        Args:
            x (torch.Tensor): The input tensor.

        Returns:
            torch.Tensor: The normalized tensor.

        """
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)

    def forward(self, x):
        """
        Forward pass through the RMSNorm layer.

        Args:
            x (torch.Tensor): The input tensor.

        Returns:
            torch.Tensor: The output tensor after applying RMSNorm.


Complete the function implementation replacing "raise NotImplementedError".