Drama Base

🚀 DRAMA-base (0.1B)：从大语言模型到小型密集检索器的多样化增强

DRAMA-base (0.1B) 是一个基于剪枝大语言模型骨干构建的密集检索模型。它通过对大语言模型进行剪枝得到，并针对高效且具有泛化能力的多语言文本检索进行了微调。 通过利用大语言模型进行高质量的数据增强，尽管 DRAMA-base 仅有 0.1B 非嵌入参数的紧凑规模，但它在英语和多语言检索任务中均取得了出色的性能。

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

更多详细信息请查看我们的 论文。

🚀 快速开始

✨ 主要特性

• 多语言支持：支持阿拉伯语、孟加拉语、中文、英语等 20 种语言，在多语言检索任务中表现出色。
• 高效紧凑：基于剪枝大语言模型骨干，仅有 0.1B 非嵌入参数，却能实现高效的文本检索。
• 灵活维度：采用套娃表示学习，嵌入维度可灵活截断。

📦 安装指南

文档未提及安装步骤，可参考 transformers 库的安装方式。

💻 使用示例

基础用法

以下是使用 drama-base 对 MIRACL 数据集中的查询和文档示例进行编码的示例，分别展示了使用 Transformers 和 Sentence Transformers 的方法：

Transformers

import torch
from transformers import AutoTokenizer, AutoModel

queries = [
    'What percentage of the Earth\'s atmosphere is oxygen?',
    '意大利首都是哪里？',
]
documents = [
    "The amount of oxygen in the atmosphere has fluctuated over the last 600 million years, reaching a peak of 35% during the Carboniferous period, significantly higher than today's 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())
# Expected output: [[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())
# Expected output: [[0.6031, 0.1750], [0.2005, 0.7251]]

Sentence Transformers

from sentence_transformers import SentenceTransformer

queries = [
    'What percentage of the Earth\'s atmosphere is oxygen?',
    '意大利首都是哪里？',
]
documents = [
    "The amount of oxygen in the atmosphere has fluctuated over the last 600 million years, reaching a peak of 35% during the Carboniferous period, significantly higher than today's 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())
# Expected output: [[0.5310, 0.0821], [0.1298, 0.6181]]

⚠️ 重要提示

• trust_remote_code 将使用我们自定义的 drama_modeling.py，该文件使用双向注意力而不是单向注意力。
• 对于查询，你必须使用 prompt_name="query" 来选择 名为 "query" 的提示，或者使用 prompt="Query: " 手动指定提示字符串。

高级用法

from sentence_transformers import SentenceTransformer

queries = [
    'What percentage of the Earth\'s atmosphere is oxygen?',
    '意大利首都是哪里？',
]
documents = [
    "The amount of oxygen in the atmosphere has fluctuated over the last 600 million years, reaching a peak of 35% during the Carboniferous period, significantly higher than today's 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())
# Expected output: [[0.6031, 0.1750], [0.2005, 0.7251]]

📚 详细文档

评估

该模型已在多个检索基准上进行了评估，包括 BEIR、MIRACL、MLDR 以及 MTEB 中的多个多语言检索任务。 它在英语和多语言检索任务中均表现出色。

本页面发布的 drama-base 对应于具有 1.13 亿非嵌入参数的 DRAMA - 0.1B 这一行。

支持的语言

DRAMA-base 从 Llama3.2 - 1B 初始化（该模型最初从 Llama3.1 - 8B 剪枝而来）。在剪枝和检索器训练期间，训练数据涵盖了以下 20 种语言（按字母顺序排列）：

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

其他语言的性能可能会下降。

📄 许可证

本项目采用 CC - BY - NC - 4.0 许可证。

📚 引用

如果您发现我们的论文或模型有帮助，请考虑按以下方式引用：

@article{drama,
  title={{Drama}: Diverse Augmentation from Large Language Models To Smaller Dense Retrievers},
  author={Ma, Xueguang and Lin, Victoria Xi and Oguz, Barlas and Lin, Jimmy and Yih, Wen-tau and Chen, Xilun},
  journal={arXiv:2502.18460},
  year={2025}
}