# -*- 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)