license: cc-by-nc-4.0 language:

• 阿拉伯语
• 孟加拉语
• 中文
• 英语
• 芬兰语
• 法语
• 德语
• 印地语
• 印尼语
• 意大利语
• 日语
• 韩语
• 波斯语
• 葡萄牙语
• 俄语
• 西班牙语
• 斯瓦希里语
• 泰卢固语
• 泰语
• 约鲁巴语 pipeline_tag: 句子相似度 library_name: transformers tags:
• 句子转换器

DRAMA-base (0.1B)：基于大语言模型增强的小型稠密检索模型

DRAMA-base (0.1B) 是一个基于剪枝后大语言模型架构的稠密检索模型。该模型通过对大语言模型进行剪枝，并针对高效且可泛化的多语言文本检索任务进行微调而得到。通过利用大语言模型进行高质量数据增强，尽管其非嵌入参数仅有0.1B的紧凑规模，DRAMA-base在英语和多语言检索任务中均展现出强劲性能。

drama-base的默认嵌入维度为768。由于采用套娃表示学习（Matryoshka Representation Learning），其维度可灵活截断至512或256等不同尺寸。

详情请参阅我们的论文。

使用方式

以下示例展示如何使用drama-base对MIRACL数据集中的查询和文档进行编码，支持Transformers和Sentence Transformers两种方式：

Transformers

import torch
from transformers import AutoTokenizer, AutoModel


queries = [
    '地球大气中氧气的百分比是多少？',
    '意大利首都是哪里？',
]
documents = [
    "过去6亿年间大气中的氧气含量持续波动，在石炭纪时期达到35%的峰值，显著高于现今21%的水平。",
    "罗马是欧洲国家意大利的首都，也是罗马广域市的行政中心，同时作为意大利全国政治、经济、文化和交通枢纽，坐落于亚平宁半岛中部的台伯河下游平原。城市始建于七座山丘之上，故有'七丘之城'的别称。按城区人口统计，罗马不仅是意大利人口最多的城市，也是欧盟人口第三大城市。",
]

model_name = "facebook/drama-base"
device = "cuda" if torch.cuda.is_available() else "cpu"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name, trust_remote_code=True).to(device)

query_embs = model.encode_queries(tokenizer, queries)
doc_embs = model.encode_documents(tokenizer, documents)

scores = query_embs @ doc_embs.T
print(scores.tolist())
# 预期输出：[[0.5310, 0.0821], [0.1298, 0.6181]]

设置trust_remote_code将启用我们定制的drama_modeling.py，包含两个特性：

• 使用双向注意力机制而非单向注意力
• 自动为查询文本添加"Query: "前缀（文档不添加前缀）

DRAMA模型采用套娃表示学习（MRL）训练，支持灵活维度调整。以下示例展示如何将查询和文档编码为256维：

query_embs = model.encode_queries(tokenizer, queries, dim=256)
doc_embs = model.encode_documents(tokenizer, documents, dim=256)

scores = query_embs @ doc_embs.T
print(scores.tolist())
# 预期输出：[[0.6031, 0.1750], [0.2005, 0.7251]]

Sentence Transformers

from sentence_transformers import SentenceTransformer

queries = [
    '地球大气中氧气的百分比是多少？',
    '意大利首都是哪里？',
]
documents = [
    "过去6亿年间大气中的氧气含量持续波动，在石炭纪时期达到35%的峰值，显著高于现今21%的水平。",
    "罗马是欧洲国家意大利的首都，也是罗马广域市的行政中心，同时作为意大利全国政治、经济、文化和交通枢纽，坐落于亚平宁半岛中部的台伯河下游平原。城市始建于七座山丘之上，故有'七丘之城'的别称。按城区人口统计，罗马不仅是意大利人口最多的城市，也是欧盟人口第三大城市。",
]

model = SentenceTransformer("facebook/drama-base", trust_remote_code=True)

query_embs = model.encode(queries, prompt_name="query")
doc_embs = model.encode(documents)

scores = model.similarity(query_embs, doc_embs)
print(scores.tolist())
# 预期输出：[[0.5310, 0.0821], [0.1298, 0.6181]]

• 设置trust_remote_code将启用定制化的drama_modeling.py，使用双向注意力机制
• 对于查询编码，必须使用prompt_name="query"选择配置文件中的"query"提示词，或手动指定prompt="Query: "

以下展示256维编码示例：

from sentence_transformers import SentenceTransformer

queries = [
    '地球大气中氧气的百分比是多少？',
    '意大利首都是哪里？',
]
documents = [
    "过去6亿年间大气中的氧气含量持续波动，在石炭纪时期达到35%的峰值，显著高于现今21%的水平。",
    "罗马是欧洲国家意大利的首都，也是罗马广域市的行政中心，同时作为意大利全国政治、经济、文化和交通枢纽，坐落于亚平宁半岛中部的台伯河下游平原。城市始建于七座山丘之上，故有'七丘之城'的别称。按城区人口统计，罗马不仅是意大利人口最多的城市，也是欧盟人口第三大城市。",
]

model = SentenceTransformer("facebook/drama-base", truncate_dim=256, trust_remote_code=True)

query_embs = model.encode(queries, prompt_name="query")
doc_embs = model.encode(documents)

scores = model.similarity(query_embs, doc_embs)
print(scores.tolist())
# 预期输出：[[0.6031, 0.1750], [0.2005, 0.7251]]

评估结果

该模型已在BEIR、MIRACL、MLDR及MTEB的多个多语言检索任务上进行评估。在英语和多语言检索任务中均表现优异。

本页发布的drama-base对应评测中的DRAMA-0.1B版本（1.13亿非嵌入参数）。

支持语言

DRAMA-base基于Llama3.2-1B（该模型本身是从Llama3.1-8B剪枝得到）初始化。在剪枝和检索器训练过程中，训练数据涵盖以下20种语言（按字母排序）：

阿拉伯语、孟加拉语、中文、英语、芬兰语、法语、德语、印地语、印尼语、意大利语、日语、韩语、波斯语、葡萄牙语、俄语、西班牙语、斯瓦希里语、泰卢固语、泰语、约鲁巴语

其他语言的性能可能有所降低。

引用

如果您觉得我们的论文或模型有帮助，请参考以下引用格式：

@article{drama,
  title={{Drama}: 基于大语言模型增强的小型稠密检索模型},
  author={马学广 and 林维多利亚 and 欧古兹·巴拉斯 and 林吉米 and 易文涛 and 陈希伦},
  journal={arXiv:2502.18460},
  year={2025}
}