model/inference.py

# -*- coding: utf-8 -*-
import pandas as pd
import numpy as np
import os
import random
from tqdm import tqdm
from sklearn.metrics import auc, roc_auc_score, precision_recall_curve
from sklearn import metrics
from torch.utils.data import DataLoader
import torch



def Seed_everything(seed=2024):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)
    if hasattr(torch.backends, 'cudnn'):
        torch.backends.cudnn.deterministic = True


def Metric(preds, labels,best_threshold = None):
    labels = np.array(labels).reshape(-1)
    preds = np.array(preds).reshape(-1)
    if best_threshold == None:
        best_f1 = 0
        best_threshold = 0
        for threshold in range(0, 100):
            threshold = threshold / 100
            binary_pred = [1 if pred >= threshold else 0 for pred in preds]
            binary_true = labels
            f1 = metrics.f1_score(binary_true, binary_pred)
            if f1 > best_f1:
                best_f1 = f1
                best_threshold = threshold
    binary_pred = [1 if pred >= best_threshold else 0 for pred in preds]
    binary_true = labels
    binary_acc = metrics.accuracy_score(binary_true, binary_pred)
    precision = metrics.precision_score(binary_true, binary_pred)
    recall = metrics.recall_score(binary_true, binary_pred)
    f1 = metrics.f1_score(binary_true, binary_pred)
    mcc = metrics.matthews_corrcoef(binary_true, binary_pred)
    AUC = roc_auc_score(labels, preds)
    precisions, recalls, _ = precision_recall_curve(labels, preds)  #######
    AUPRC = auc(recalls, precisions)
    return AUC, AUPRC, mcc ,binary_acc,precision,recall,f1


def Write_log(logFile, text, isPrint=True):
    if isPrint:
        print(text)
    logFile.write(text)
    logFile.write('\n')
    return None


class TaskDataset:
    def __init__(self, df, protein_data, label_name):
        self.df = df
        self.protein_data = protein_data
        self.label_name = label_name

    def __len__(self):
        return (self.df.shape[0])

    def __getitem__(self, idx):
        pdb_id = self.df.loc[idx, 'ID']

        protein_X, protein_node_features, protein_masks, labels, adj = self.protein_data[pdb_id]

        return {
            'PDB_ID': pdb_id,
            'PROTEIN_X': protein_X,
            'PROTEIN_NODE_FEAT': protein_node_features,
            'PROTEIN_MASK': protein_masks,
            'LABEL': labels,
            'ADJ': adj,
        }

    def collate_fn(self, batch):
        pdb_ids = [item['PDB_ID'] for item in batch]
        protein_X = torch.stack([item['PROTEIN_X'] for item in batch], dim=0)
        protein_node_features = torch.stack([item['PROTEIN_NODE_FEAT'] for item in batch], dim=0)
        protein_masks = torch.stack([item['PROTEIN_MASK'] for item in batch], dim=0)
        labels = torch.stack([item['LABEL'] for item in batch], dim=0)
        adj = torch.stack([item['ADJ'] for item in batch], dim=0)

        return pdb_ids, protein_X, protein_node_features, protein_masks, labels, adj


# main function
def model_test(test, protein_data, model_class, config, weight_dir, output_dir,
               logit=False, device=None, num_workers=0):
    label_name = ['label']  # some task may have mutiple labels
    sequence_name = "sequence"
    device = torch.device(device or ('cuda' if torch.cuda.is_available() else 'cpu'))
    print("Device:", device)

    output_result = os.fspath(output_dir)
    output_weight = os.fspath(weight_dir)
    os.makedirs(output_result, exist_ok=True)


    node_features = config['node_features']
    edge_features = config['edge_features']
    hidden_dim = config['hidden_dim']
    num_encoder_layers = config['num_encoder_layers']
    k_neighbors = config['k_neighbors']
    augment_eps = config['augment_eps']
    dropout = config['dropout']
    id_name = config['id_name']
    batch_size = config['batch_size']
    folds = config['folds']


    if test is not None:

        log = open(os.path.join(output_result, 'test.log'), 'w', buffering=1)
        Write_log(log, str(config) + '\n')
        sub = test[[id_name, sequence_name]].copy()

        if isinstance(label_name, list):
            for l in label_name:
                sub[l] = 0.0
                sub[l] = sub[l].astype(np.float32)
        else:
            sub[label_name] = 0.0

        test_dataset = TaskDataset(test, protein_data, label_name)
        loader_kwargs = dict(
            dataset=test_dataset,
            batch_size=batch_size,
            collate_fn=test_dataset.collate_fn,
            shuffle=False,
            drop_last=False,
            num_workers=num_workers,
        )
        if num_workers > 0:
            loader_kwargs['prefetch_factor'] = 2
        test_dataloader = DataLoader(**loader_kwargs)
        models = []
        for fold in range(folds):
            checkpoint = os.path.join(output_weight, 'fold%s.ckpt' % fold)
            if not os.path.exists(checkpoint):
                print("Missing checkpoint:", checkpoint)
                continue

            model = model_class(node_features, edge_features, hidden_dim, num_encoder_layers, k_neighbors, augment_eps, dropout)
            model.to(device)
            state_dict = torch.load(checkpoint, map_location=device, weights_only=True)
            model.load_state_dict(state_dict)

            model.eval()
            models.append(model)
        print('model count:', len(models))
        if not models:
            raise FileNotFoundError(f'No fold*.ckpt files were loaded from {output_weight}')

        test_preds = []
        test_outputs = []
        test_Y = []
        all_protein_node_features = []
        all_labels = []
        with torch.no_grad():
            for data in tqdm(test_dataloader):
                protein_X, protein_node_features, protein_masks, y, adj = [d.to(device) for d in data[1:]]
                all_protein_node_features.append(protein_node_features.detach().cpu().numpy())
                all_labels.append(y.detach().cpu().numpy())
                if logit:
                    outputs = [model(protein_X, protein_node_features, protein_masks, adj).sigmoid() for model in models]
                else:
                    outputs = [model(protein_X, protein_node_features, protein_masks) for model in models]

                outputs = torch.stack(outputs, 0).mean(0)  # 5个模型预测结果求平均,最终shape=(bsize, max_len)
                test_outputs.append(outputs.detach().cpu().numpy())

                test_seq_y = torch.masked_select(y, protein_masks.bool())
                test_seq_preds = torch.masked_select(outputs, protein_masks.bool())

                test_preds.append(test_seq_preds.cpu().detach().numpy())
                test_Y.append(test_seq_y.cpu().detach().numpy())


        test_preds = np.concatenate(test_preds)
        test_Y = np.concatenate(test_Y)
        test_metric = Metric(test_preds, test_Y)
        Write_log(log,'test_auc:%.6f, test_auprc:%.6f, testFYT_mccL:%.6f, test_acc:%.6f, test_pre:%.6f, test_rec:%.6f, test_f1:%.6f' \
                      % (test_metric[0], test_metric[1], test_metric[2], test_metric[3],
                         test_metric[4], test_metric[5], test_metric[6]))

        test_outputs = np.concatenate(test_outputs)  # shape = (num_samples, max_len) or (num_samples,  4 * max_len)


        sub['label'] = sub['label'].astype(object)
        for i in range(len(sub)):
            sub.at[i, 'label'] = test_outputs[i, :len(sub.loc[i, sequence_name])].tolist()
        sub.to_csv(os.path.join(output_result, 'result.csv'), index=False)
        log.close()