语言:

• 英语 标签:
• 特征提取
• 生物医学文献
• 句子相似度 数据集:
• biu-nlp/abstract-sim-pubmed

一个用于将摘要句子描述映射到符合描述的句子的模型。该模型基于生物医学文献句子进行训练。使用load_finetuned_model加载查询和句子编码器，使用encode_batch()通过模型对句子进行编码。

from transformers import AutoTokenizer, AutoModel
import torch

def load_finetuned_model():
        sentence_encoder = AutoModel.from_pretrained("biu-nlp/abstract-sim-sentence-pubmed", revision="71f4539120e29024adc618173a1ed5fd230ac249")
        query_encoder = AutoModel.from_pretrained("biu-nlp/abstract-sim-query-pubmed", revision="8d34676d80a39bcbc5a1d2eec13e6f8078496215")
        tokenizer = AutoTokenizer.from_pretrained("biu-nlp/abstract-sim-sentence-pubmed")
        return tokenizer, query_encoder, sentence_encoder


def encode_batch(model, tokenizer, sentences, device):
    input_ids = tokenizer(sentences, padding=True, max_length=128, truncation=True, return_tensors="pt",
                          add_special_tokens=True).to(device)
    features = model(**input_ids)[0]

    features =  torch.sum(features[:,:,:] * input_ids["attention_mask"][:,:].unsqueeze(-1), dim=1) / torch.clamp(torch.sum(input_ids["attention_mask"][:,:], dim=1, keepdims=True), min=1e-9)

    return features