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}
}