标签:

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

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