标签:

• 特征提取 管道标签: 特征提取

DRAGON+ 是一个基于BERT架构的密集检索模型，初始权重源自RetroMAE，并采用《如何训练你的DRAGON：面向通用密集检索的多样化数据增强》中描述的方法，在MS MARCO语料库增强数据上进行了进一步训练。

相关GitHub仓库详见 https://github.com/facebookresearch/dpr-scale/tree/main/dragon。我们采用非对称双编码器架构，配备两个独立参数化的编码器。其他可用模型如下：

使用方式（HuggingFace Transformers）

直接调用HuggingFace transformers库中的模型：

import torch
from transformers import AutoTokenizer, AutoModel
tokenizer = AutoTokenizer.from_pretrained('facebook/dragon-plus-query-encoder')
query_encoder = AutoModel.from_pretrained('facebook/dragon-plus-query-encoder')
context_encoder = AutoModel.from_pretrained('facebook/dragon-plus-context-encoder')

# 以MS MARCO查询和段落为例
query =  "玛丽·居里出生在哪里？"
contexts = [
    "玛丽亚·斯克沃多夫斯卡，后被称为玛丽·居里，于1867年11月7日出生。",
    "皮埃尔·居里1859年5月15日出生于巴黎，其父尤金·居里是来自阿尔萨斯地区的法裔天主教医生。"
]
# 应用分词器
query_input = tokenizer(query, return_tensors='pt')
ctx_input = tokenizer(contexts, padding=True, truncation=True, return_tensors='pt')
# 计算嵌入向量：获取[CLS]标记的最后一层隐藏状态
query_emb = query_encoder(**query_input).last_hidden_state[:, 0, :]
ctx_emb = context_encoder(**ctx_input).last_hidden_state[:, 0, :]
# 使用点积计算相似度分数
score1 = query_emb @ ctx_emb[0]  # 396.5625
score2 = query_emb @ ctx_emb[1]  # 393.8340