许可证: mit 数据集:

• MedRAG/教材
• MedRAG/PubMed医学文献
• MedRAG/StatPearls医学百科
• mteb/生物预印本原始数据
• mteb/医学预印本原始数据
• MS MARCO问答数据集
• BMRetriever/生物医学检索数据集 支持语言:
• 英语 标签:
• 医学
• 生物学
• 信息检索
• 大语言模型

本模型基于论文《BMRetriever：将大语言模型调优为更优的生物医学文本检索器》所述方法进行微调，相关GitHub代码库详见 https://github.com/ritaranx/BMRetriever。

该模型参数量达70亿，技术细节请参阅论文链接。

使用指南

通过HuggingFace transformers库加载预训练模型：

from transformers import AutoModel, AutoTokenizer

model = AutoModel.from_pretrained("BMRetriever/BMRetriever-7B") 
tokenizer = AutoTokenizer.from_pretrained("BMRetriever/BMRetriever-7B") 

获取句子嵌入向量的方法如下：

import torch
import torch.nn.functional as F

from torch import Tensor
from transformers import AutoTokenizer, AutoModel


def 末词池化(last_hidden_states: Tensor,
                 attention_mask: Tensor) -> Tensor:
    last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
    left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
    if left_padding:
        embedding = last_hidden[:, -1]
    else:
        sequence_lengths = attention_mask.sum(dim=1) - 1
        batch_size = last_hidden.shape[0]
        embedding = last_hidden[torch.arange(batch_size, device=last_hidden.device), sequence_lengths]
    return embedding

def 生成详细指令查询(task_description: str, query: str) -> str:
    return f'{task_description}\n查询：{query}'

def 生成详细指令文本(passage: str) -> str:
    return f'表征以下文本\n文本：{passage}'

# 每个查询需附带任务描述
task = '给定科学论断，检索支持或反驳该论断的文献'
queries = [
    生成详细指令查询(task, '顺式作用长链非编码RNA可调控其转录位点附近基因的表达'),
    生成详细指令查询(task, '叉头框0（Fox0）转录因子参与细胞凋亡过程')
]

# 检索文档无需添加指令
documents = [
    生成详细指令文本("长链非编码RNA转录行为对基因的调控作用：长链非编码RNA（lncRNA）被认为是小鼠和人类转录组中最大的转录本类别。两个关键问题在于是否所有lncRNA都具有功能，以及它们如何发挥作用。已有研究表明部分lncRNA通过其产物发挥作用，但这并非唯一作用模式。本文重点探讨lncRNA转录过程（独立于其产物）对邻近蛋白质编码基因活性的顺式调控作用，包括导致基因沉默或激活的实例，并介绍区分转录过程与RNA产物作用机制的研究策略。"),
    生成详细指令文本("增强子非编码转录：基本原理与功能模型。哺乳动物基因组在蛋白质编码基因边界外存在广泛转录现象。近期全基因组研究表明，增强子和基因座控制区等转录调控元件是非编码转录的主要位点。虽然增强子来源的RNA在细胞非核糖体RNA中占比较小，但越来越多实验数据表明增强子转录及其产物eRNA可能具有特定功能，而非仅是开放染色质区域的转录噪声。本文综述当前关于增强子转录及其功能意义的研究进展。")
]
input_texts = queries + documents

max_length = 512

# 文本向量化处理
batch_dict = tokenizer(input_texts, max_length=max_length-1, padding=True, truncation=True, return_tensors='pt')

# 关键步骤！在末尾添加EOS标记
batch_dict['input_ids'] = [input_ids + [tokenizer.eos_token_id] for input_ids in batch_dict['input_ids']]
batch_dict = tokenizer.pad(batch_dict, padding=True, return_attention_mask=True, return_tensors='pt').to("cuda")

model.eval()
with torch.no_grad():
    outputs = model(**batch_dict)
    embeddings = 末词池化(outputs.last_hidden_state, batch_dict['attention_mask'])

通过嵌入向量的点积计算句子相似度：

scores = (embeddings[:2] @ embeddings[2:].T)
print(scores.tolist())

引用规范

若本项目对您的研究有所帮助，请引用以下论文：

@misc{xu2024bmretriever,
      title={BMRetriever：将大语言模型调优为更优的生物医学文本检索器}, 
      author={徐然 and 施文琪 and 余越 and 庄雨辰 and 朱彦桥 and 王美冬 and 何静怡 and 张超 and 杨卡尔},
      year={2024},
      eprint={2404.18443},
      archivePrefix={arXiv},
      primaryClass={cs.CL}
}