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