a_predict222.py

# -*- conding: utf-8 -*-
# @Time    : 2025/12/14  10:58
# @Author  : psi

from utils import *
from modules import *
import os, sys
import numpy as np
from tqdm import tqdm
import random
import torch
from torch import nn
from config import CFG
from dataset import *
import torch.utils.data
import copy, json, pickle
import itertools as it
import glob
import torch.nn.functional as F


def my_collate(batch):
    batch = list(filter(lambda x: (x is not None), batch))
    msbinl, molfpl, molfml, vl, al, msl = [], [], [], [], [], []
    bat = {}
    msbinl1, msbinl2 = [], []

    for b in batch:
        if 'ms_bins' in b:
            msbinl.append(b['ms_bins'])
        if 'ms_bins1' in b:
            msbinl1.append(b['ms_bins1'])
        if 'ms_bins2' in b:
            msbinl2.append(b['ms_bins2'])
        if 'mol_fps' in b:
            molfpl.append(b['mol_fps'])
        if 'mol_fmvec' in b:
            molfml.append(b['mol_fmvec'])
        if 'V' in b:
            vl.append(b['V'])
        if 'A' in b:
            al.append(b['A'])
        if 'mol_size' in b:
            msl.append(b['mol_size'])

    if msbinl:
        bat['ms_bins'] = torch.stack(msbinl)
    if msbinl1:
        bat['ms_bins1'] = torch.stack(msbinl1)
    if msbinl2:
        bat['ms_bins2'] = torch.stack(msbinl2)
    if molfpl:
        bat['mol_fps'] = torch.stack(molfpl)
    if molfml:
        bat['mol_fmvec'] = torch.stack(molfml)
    if vl and al and msl:
        max_n = max(map(lambda x:x.shape[0], vl))
        vl1, al1 = [], []
        for v in vl:
            vl1.append(pad_V(v, max_n))
        for a in al:
            al1.append(pad_A(a, max_n))

        bat['V'] = torch.stack(vl1)
        bat['A'] = torch.stack(al1)
        bat['mol_size'] = torch.cat(msl, dim=0)

    # return torch.utils.data.dataloader.default_collate(batch)
    return bat


def build_loaders(inp, mode, cfg, num_workers):
    if type(inp[0]) is dict:
        dataset = Dataset(inp, cfg)
    else:
        dataset = PathDataset(inp, cfg)
    dataloader = torch.utils.data.DataLoader(
        dataset,
        batch_size=len(dataset),
        num_workers=num_workers,
        shuffle=True if mode == "train" else False,
        collate_fn=my_collate
    )
    return dataloader


class Predictor():

    def __init__(self, file, model_file):
        CFG.load(file)
        cfg = CFG
        self.cfg = cfg
        model = FragSimiModelNew(cfg).to(cfg.device)
        encmodel = torch.load(model_file)
        # model.mol_gnn_encoder.load_state_dict(encmodel.mol_gnn_encoder.state_dict())
        model.load_state_dict(encmodel['state_dict'])

        self.model = model
        self.model.eval()
        self.device = "cuda" if torch.cuda.is_available() else "cpu"

    def process_file(self, ms, smi):
        # d = json.load(open(file, 'r', encoding='utf-8'))
        # ms = d['ms']
        # smi = d['smiles']
        # out = {'ms': ms, 'smiles': smi}
        # ms = self.data[idx]['ms']
        # smi = self.data[idx]['smiles']
        nls = []
        item = calc_feats(smi, ms, nls, self.cfg)
        return item

    def process(self, data):

        res = []

        for d in data:
            try:
                res.append([d['smiles'], self.process_file(d['ms'], d['smiles'])])
            except Exception as e:
                print(e)

        if len(res) == 0:
            return None

        res1 = [x[1] for x in res]
        res2 = [x[0] for x in res]

        batch = my_collate(res1)

        return batch, res2

    def get_eval_info(self, ms_embeddings, mol_embeddings, top_ks=(1, 3, 5, 10)):
        N = ms_embeddings.shape[0]

        # 1. L2 归一化（非常关键）
        # ms_norm = F.normalize(ms_embeddings, dim=1)
        # mol_norm = F.normalize(mol_embeddings, dim=1)
        ms_norm = ms_embeddings
        mol_norm = mol_embeddings

        recalls = {k: 0 for k in top_ks}

        # 2. 对每个样本做检索
        for i in range(N):
            query = ms_norm[i]  # (256,)
            sims = torch.matmul(mol_norm, query)  # (N,)

            ranked_indices = torch.argsort(sims, descending=True)

            for k in top_ks:
                if i in ranked_indices[:k]:
                    recalls[k] += 1

        # 3. 取平均
        for k in recalls:
            recalls[k] /= N

        return recalls

    def topk_similarity(self, ms_embedding, res_embeddings, batch_size=128, top_k=10):

        ms_embedding = ms_embedding.to(self.device)
        res_embeddings = res_embeddings.to(self.device)

        # 保证是 float tensor
        ms_embedding = ms_embedding.float()
        res_embeddings = res_embeddings.float()

        # 归一化（用于余弦相似度）
        # ms_embedding = F.normalize(ms_embedding, dim=1)
        # res_embeddings = F.normalize(res_embeddings, dim=1)

        similarities = []

        # 分 batch 计算
        for i in range(0, res_embeddings.size(0), batch_size):
            batch = res_embeddings[i:i + batch_size]  # (B, 256)

            # (1, 256) @ (256, B) -> (1, B)
            sim = torch.matmul(ms_embedding, batch.T)  # 余弦相似度
            similarities.append(sim.squeeze(0))  # (B,)

        # 拼接成 (15511,)
        similarities = torch.cat(similarities, dim=0)

        top_k = min(top_k, res_embeddings.size(0))
        # 取 top_k
        topk_sim, topk_idx = torch.topk(similarities, k=top_k)

        return topk_sim, topk_idx

    def predict(self, data):
        # data_files = []
        # for root, _, files in os.walk(file_path):
        #     for f in files:
        #         if f.endswith(('.json', '.pkl', '.mgf')):
        #             data_files.append(os.path.join(root, f))
        # data = sorted(
        #     data_files,
        #     key=lambda x: int(os.path.splitext(os.path.basename(x))[0])
        # )
        batch, res_smiles = self.process(data)

        if batch is None:
            return None

        for k, v in batch.items():
            batch[k] = v.to(self.cfg.device)

        loss, loss_infonce, loss_mse, ms_embeddings, mol_embeddings = self.model(batch, is_predict=True)

        # recalls_info = self.get_eval_info(ms_embeddings, mol_embeddings)
        # print(recalls_info)
        #
        # return loss, loss_infonce, loss_mse, recalls_info

        ms_embedding = ms_embeddings[0:1, :]

        topk_sim, topk_idx = self.topk_similarity(ms_embedding, mol_embeddings)
        topk_idx = topk_idx.to("cpu").numpy().tolist()

        res_pred_name = []
        for x, i in enumerate(topk_idx):
            res_pred_name.append([res_smiles[i], topk_sim[x].item()])

        return res_pred_name


model_file = ["/root/代码/out_data/train-020/model-tloss3.239-vloss2.752-epoch0.pth",
              "/root/代码/out_data/train-020/model-tloss2.487-vloss2.279-epoch1.pth",
              "/root/代码/out_data/train-020/model-tloss2.086-vloss1.924-epoch2.pth",
              "/root/代码/out_data/train-020/model-tloss1.716-vloss1.609-epoch3.pth",
              "/root/代码/out_data/train-020/model-tloss1.414-vloss1.361-epoch4.pth",
              '/root/代码/out_data/train-020/model-tloss1.177-vloss1.173-epoch5.pth',
              "/root/代码/out_data/train-020/model-tloss0.99-vloss1.036-epoch6.pth",
              '/root/代码/out_data/train-020/model-tloss0.849-vloss0.929-epoch7.pth',
              '/root/代码/out_data/train-020/model-tloss0.736-vloss0.851-epoch8.pth',
              '/root/代码/out_data/train-020/model-tloss0.647-vloss0.779-epoch9.pth',
              '/root/代码/out_data/train-020/model-tloss0.575-vloss0.728-epoch10.pth'][-1]

pred = Predictor('config.json', model_file)


if __name__ == '__main__':


    # file_path = ["../../data/CASMI2016/data/neg", "../../data/CASMI2017/data/neg"][0]
    ms = [
        [
            107.049103,
            11496.099609
        ],
        [
            131.050201,
            12601.5
        ],
        [
            157.053299,
            11285.700195
        ],
        [
            158.061096,
            2146837.0
        ],
        [
            222.022995,
            10585360.0
        ]
    ]
    data = [{'ms': ms, 'smiles': "NC1=C2C=CC(=CC2=CC=C1)S(O)(=O)=O"},
            {'ms': ms, 'smiles': "OS(=O)(=O)c1ccc2ccc(cc2c1)S(O)(=O)=O"},
            {'ms': ms, 'smiles': "OC1=CC=CC2=C(C=CC=C12)S(O)(=O)=O"},
            {'ms': ms, 'smiles': "OC(=O)c1nc(Cl)ccc1Cl"},
            {'ms': ms, 'smiles': "CN1C=NC(=C1SC1=C2NC=NC2=NC=N1)[N+]([O-])=O"},
            {'ms': ms, 'smiles': "CC(=O)NC1=CC=C(O)C(=C1)C(O)=O"}]
    res_pred_name = pred.predict(data)

    print(res_pred_name)