语言:

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

该模型用于将摘要句子描述映射到符合描述的句子。基于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