embedding_predict111.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()

    def process_file(self, smi):
        # d = json.load(open(file, 'r', encoding='utf-8'))
        # ms = d['ms']
        # smi = d['smiles']
        ms = [[41.038587, 880600.0], [42.033833, 1973400.0], [43.041651, 2117400.0], [44.049388, 925150.0], [44.979347, 4397200.0], [51.022884, 593400.0], [53.038537, 9694400.0], [54.033783, 415000.0], [55.054152, 1911200.0], [56.049325, 4400500.0], [56.979301, 487200.0], [65.038474, 449400.0], [67.041567, 1667200.0], [67.054107, 786000.0], [68.049268, 6593250.0], [68.979253, 836000.0], [69.056975, 17628050.0], [69.069744, 290600.0], [70.06482, 1716900.0], [70.994926, 276400.0], [73.010627, 215600.0], [77.038437, 417800.0], [79.054112, 1319600.0], [80.049378, 3584000.0], [81.057194, 2957200.0], [82.064812, 25688800.0], [82.070396, 669000.0], [82.073249, 528650.0], [82.994909, 3564600.0], [83.072653, 7343000.0], [84.080502, 821400.0], [94.065026, 1006000.0], [95.049053, 230800.0], [96.080647, 13938800.0], [97.010531, 20776600.0], [97.013298, 339600.0], [98.989853, 367600.0], [110.096244, 1418800.0], [110.989833, 48727600.0], [110.991981, 1024000.0], [110.994248, 515600.0], [111.001067, 985400.0], [111.103933, 13806250.0], [112.111972, 17873400.0], [112.114998, 263400.0], [115.054168, 518400.0], [117.069717, 320200.0], [134.018401, 474400.0], [194.099834, 1533000.0], [194.993274, 21076400.0], [306.09803, 51809350.0], [306.181335, 516550.0]]

        # 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, file):
        if isinstance(file, str):
            res = json.load(open(file, 'r', encoding='utf-8'))
            data = res['smiles']
            res = []
            for d in tqdm(data, desc='process ...'):
                try:
                    res.append([d, self.process_file(d)])
                except:
                    res.append([d, None])

            o_file = '/dev/shm/data/tongji_data/all_pos_pred.pt'
            torch.save(res, o_file)

            os._exit(0)

        else:
            res = file
            batch = my_collate(res)
            return batch

    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 predict(self, file_path):
        # 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 = self.process(file_path)
        for k, v in batch.items():
            batch[k] = v.to(self.cfg.device)

        with torch.no_grad():
            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

        return mol_embeddings


if __name__ == '__main__':
    model_file = ["model-tloss3.437-vloss2.907-epoch0.pth", "model-tloss2.495-vloss2.253-epoch1.pth",
                  "model-tloss1.987-vloss1.866-epoch2.pth", "model-tloss1.597-vloss1.573-epoch3.pth",
                  "model-tloss1.332-vloss1.384-epoch4.pth", "model-tloss1.088-vloss1.255-epoch5.pth",
                  "model-tloss0.899-vloss1.068-epoch6.pth",
                  "/root/代码/out_data/train-018/model-tloss0.76-vloss0.986-epoch0.pth",
                  "/root/代码/out_data/train-019/model-tloss0.608-vloss0.943-epoch0.pth",
                  "/root/代码/out_data/train-019/model-tloss0.577-vloss0.852-epoch1.pth",
                  "/root/代码/out_data/train-019/model-tloss0.503-vloss0.811-epoch2.pth",
                  '/root/代码/out_data/train-019/model-tloss0.448-vloss0.763-epoch3.pth',
                  '/root/代码/out_data/train-019/model-tloss0.405-vloss0.74-epoch4.pth',
                  '/root/代码/out_data/train-019/model-tloss0.367-vloss0.722-epoch5.pth',
                  '/root/代码/out_data/train-019/model-tloss0.337-vloss0.705-epoch6.pth',
                  '/root/代码/out_data/train-019/model-tloss0.317-vloss0.671-epoch7.pth'][-1]
    model_name = model_file.split('/')[-1][:-4]
    pred = Predictor('config.json', model_file)

    # file_path = '/dev/shm/data/tongji_data/all_pos.json'
    # loss, loss_infonce, loss_mse, recalls_info = pred.predict(file_path)

    o_file = '/dev/shm/data/tongji_data/all_pos_pred.pt'
    data = torch.load(o_file)
    batch_size = 128
    res = []
    for i in tqdm(range(0, len(data), batch_size), desc='predict ...'):
        batch = data[i:i+batch_size]
        p = [x[1] for x in batch]
        mol_embeddings = pred.predict(p)
        mol_embeddings = mol_embeddings.to("cpu")
        res.append(mol_embeddings)

    result = torch.cat(res, dim=0)

    print(f"len is : {len(data)} ...")
    print(f"result shape is : {result.shape} ...")

    out_file = f'/dev/shm/data/tongji_data/all_pos_pred_emb_{model_name}.pt'
    torch.save(result, out_file)