Project: s41467-025-56173-6 | Function: Attforward (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:** Core attention-based forward propagation function that implements the multi-instance learning framework with self-attention mechanism. It processes instance features through feature extraction, applies attention mechanism, computes logistic probabilities, and aggregates results using weighted noisy-OR for m6A site prediction.

**KEY TERMS:** {
  "H": "instance embeddings h_k",
  "A": "attention weights a_k",
  "pij": "instance probabilities after logistic transformation",
  "M": "aggregated features z",
  "Y_prob": "final bag classification probability",
  "pi": "weighted noisy-OR output"
}

**CONTEXT CODE:** 
    def build_classifier(self):

        r'''
        Instance method for building classifier module according to config
        '''


        self.classifier = build_classifier(
            input=self.model_config['attention']['L'] * self.model_config['attention']['K']
        ).to(self.device)

    def build_logistic(self):

        r'''
        Instance method for building logistic module according to config
        '''

        self.A, self.B = build_logistic(device=self.device)

    def _logistic(self, loss):

        r"""
        instance method to get instance probability according to reconstruction loss

            Args:
                loss (torch.Tensor): reconstruction loss
            Return:
                (torch.Tensor): instance probability
        """
        return torch.sigmoid(self.A * loss + self.B)

    def build_model(self):

        r'''
        Instance method for building pum6a model according to config
        '''

        '1. build feature extractor'
        self.build_FE()

        '2. build attention module'
        self.build_attention()

        '3. build classifier'
        self.build_classifier()

        '4. build logistic module'
        self.build_logistic()

    def Attforward(self, x):
        # MASKED: Core attention-based forward propagation function that implements the multi-instance learning framework with self-attention mechanism. It processes instance features through feature extraction, applies attention mechanism, computes logistic probabilities, and aggregates results using weighted noisy-OR for m6A site prediction.
        raise NotImplementedError("This function is masked for testing purposes.")

    def bag_forward(self, input):

        r'''
        Instance method to get modification probability on the bag level from instance features.

               Args:
                       input (torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor):
                       Tensor representation of bag, bag_labels, and number of instance
               Returns:
                       loss (torch.Tensor): A tensor representation of model_factory loss
                       data_inst (torch.Tensor...): pi, pij, Y_prob, A
       '''


        bag, bag_labels, n_instance = input

        data_inst = [self.Attforward(item) for item in bag]
        idx_l2 = torch.where(bag_labels != 0)[0]
        if idx_l2.shape[0] > 0:
            l2 = [data_inst[item] for item in idx_l2]
            p = torch.stack([item[2] for item in l2])
            pi= torch.stack([item[0] for item in l2])
            y = torch.stack([bag_labels[index] for index in idx_l2])
            y = y.to(self.device)
            p = p.to(self.device)
            pi = pi.to(self.device)
            loss = -1*(log_diverse_density(pi, y)+1e-10) + 0.01*(self.A**2+self.B**2)[0]
            y[torch.where(y == -1)] = 0
            loss += torch.sum(-1. * (y * torch.log(p) + (1. - y) * torch.log(1. - p)))  # pro
        else:
            loss = 0.01*(self.A**2+self.B**2)[0]

        return loss, data_inst

    def validation(self, bag, bag_label):

        """
        instance method to get bag probability and likihood loss

            Args:
                bag (list): bag dataset
                bag_label (list): bag label
            Return:
                bag_loss (torch.Tensor): bag likelihood loss
        """

        data_inst = [self.Attforward(item.to(self.device)) for item in bag]
        bag_pro = torch.concat([item[0] for item in data_inst]).to(self.device)
        bag_loss = torch.sum(-1. * (bag_label * torch.log(bag_pro) + (1. - bag_label) * torch.log(1. - bag_pro)))

        return bag_loss

    def decision(self, bag):

        """
        instance method to get bag probability and instance probability

            Args:
                bag (list): bag dataset
            Return:
                bag_pro (torch.Tensor): tensor representation of bag probability
                ins_pro (torch.Tensor): tensor representation of instance probability
        """

        data_inst = [self.Attforward(item.to(self.device)) for item in bag]
        bag_pro = torch.concat([item[0] for item in data_inst]).to(self.device)
        ins_pro = torch.concat([item[1].squeeze() for item in data_inst]).to(self.device)

        return bag_pro, ins_pro

Complete the function implementation replacing "raise NotImplementedError".