ChemBERTa 化合物小分子的向量表示及相似检索

参考：https://arxiv.org/pdf/2209.01712.pdf

模型是基于分子simles进行transformer的MLM预训练的bert模型

1、化合物小分子的向量表示

from transformers import BertTokenizer, AutoTokenizer,BertModel
tokenizer_ = AutoTokenizer.from_pretrained("DeepChem/ChemBERTa-77M-MLM")

model_ = BertModel.from_pretrained("DeepChem/ChemBERTa-77M-MLM")


outputs_ = model_(**tokenizer_("Clc1c(Cl)c(O)ccc1O", return_tensors='pt'))

outputs_.pooler_output[0] ##向量，384维度

这是tokenizer_ 查看输入基构建的输入情况，添加了cls seq 两个token

这是outputs_结果，pooler_output表示cls向量，注意这里pooler_output其实是在cls上再加了层mlp的；last_hidden_state表示所有18个token的向量；这里使用的cls向量用的是last_hidden_state里的，效果对比下来会好很多

2、分子相似检索

1）计算转换分子为向量表示保存

##数据加载 1万个smi分子，文件下载可参考：https://github.com/zilliztech/MolSearch/edit/master/script/test_1w.smi

mols2 = []
with open(r"1w.smi","r") as f:
    ss = f.readlines()
    # print(ss)
    for i in ss:
        mols2.append(i.strip().split()[0])

import numpy as np
from tqdm import tqdm

mol_ids = []
vess=[]
for mol_ in tqdm(mols2):
    outputs_mol = model_(**tokenizer_(mol_, return_tensors='pt'))
    # mol_ids.append(mol_ids)
    vess.append(outputs_mol.last_hidden_state[:,0,:][0].detach().numpy())

np.save(r"1w_smiles__embs1.npy", vess)

2）检索

import torch
import numpy as np


# 加载数据
vess_all = np.load(r"1w_smiles__embs1.npy", allow_pickle=True)
vess_ = torch.Tensor([(torch.from_numpy(item) / torch.from_numpy(item).norm(dim=-1, keepdim=True)).cpu().numpy() for item in vess_all])


## 加载单个分子用来检索的
single_vec = model_(**tokenizer_(mols2[0], return_tensors='pt')).last_hidden_state[:,0,:][0].detach().numpy()
single_vec_ = torch.Tensor((torch.from_numpy(single_vec) / torch.from_numpy(single_vec).norm(dim=-1, keepdim=True)).cpu().numpy()).unsqueeze(0)

## 检索
similarities = (vess_ @ single_vec_.T).squeeze(1)

best_photo_idx = np.argsort(similarities.numpy())[::-1]

simle_ids = [mols2[i] for i in best_photo_idx[:10]]  ##取前10

3）结果展示

from rdkit.Chem import AllChem, Draw
from rdkit import Chem


Chem.MolFromSmiles(mols2[564]) ##单个分子

Draw.MolsToGridImage([ Chem.MolFromSmiles(i) for i in simle_ids],subImgSize=(300,300), molsPerRow=3)  ## 前10相似

这是mols2[564]分子的结果

Chem.MolFromSmiles(mols2[0]) ##单个分子
Draw.MolsToGridImage([ Chem.MolFromSmiles(i) for i in simle_ids],subImgSize=(300,300), molsPerRow=5)  ## 前10相似

这是mols2[0]分子的结果