SapBERT: Self-alignment pretraining for BERT的代码使用示例

最近在研究SapBERT来计算实体的相似度,发现官方的repo没有给使用示例,我仿照写了一下使用示例,方便直接把SapBERT用起来,我的环境是:

torch                   1.7.1+cu101
torchvision             0.11.3
transformers            4.16.2

下面是使用代码,知道SapBERT是抽取向量的就行了,然后就可以用一些类似faiss的近似向量检索工具进行检索了:

from transformers import AutoTokenizer, AutoModel
import numpy as np

tokenizer = AutoTokenizer.from_pretrained("cambridgeltl/SapBERT-from-PubMedBERT-fulltext")

model = AutoModel.from_pretrained("cambridgeltl/SapBERT-from-PubMedBERT-fulltext")

query = "cardiopathy"
query_toks = tokenizer.batch_encode_plus([query], 
                                       padding="max_length", 
                                       max_length=25, 
                                       truncation=True,
                                       return_tensors="pt")
print(query_toks)
query_output = model(**query_toks)
query_cls_rep = query_output[0][:,0,:]
print(query_cls_rep)

all_names = ['Neoplasm of anterior aspect of epiglottis']

toks = tokenizer.batch_encode_plus(all_names, 
                                       padding="max_length", 
                                       max_length=25, 
                                       truncation=True,
                                       return_tensors="pt")

output = model(**toks)
cls_rep = output[0][:,0,:]
print(cls_rep)

# for large-scale search, should switch to faiss
from scipy.spatial.distance import cdist

dist = cdist(query_cls_rep.cpu().detach().numpy(), cls_rep.cpu().detach().numpy())
nn_index = np.argmin(dist)
# print ("predicted label:", snomed_sf_id_pairs_100k[nn_index])

  • 0
    点赞
  • 0
    收藏
    觉得还不错? 一键收藏
  • 打赏
    打赏
  • 0
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

农民小飞侠

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值