Drama Base

🚀 DRAMA-base (0.1B): 大規模言語モデルから小型の密ベクトル検索器への多様なデータ拡張

DRAMA-base (0.1B) は、大規模言語モデルのバックボーンを剪定して構築された密ベクトル検索モデルです。大規模言語モデルを剪定し、効率的で汎用性の高い多言語テキスト検索用に微調整されています。 大規模言語モデルを活用した高品質なデータ拡張により、DRAMA-base は0.1Bの非埋め込みパラメータというコンパクトなサイズでありながら、英語および多言語の検索タスクで強力な性能を発揮します。

drama-base のデフォルトの埋め込みサイズは768です。マトリョーシカ表現学習を採用しているため、次元数を512や256などに柔軟に切り捨てることができます。

詳細については、論文をご確認ください。

🚀 クイックスタート

概要

DRAMA-base (0.1B) は、剪定された大規模言語モデルをベースに構築された密ベクトル検索モデルで、多言語のテキスト検索に最適化されています。

サポート言語

• アラビア語 (ar)
• ベンガル語 (bn)
• 中国語 (zh)
• 英語 (en)
• フィンランド語 (fi)
• フランス語 (fr)
• ドイツ語 (de)
• ヒンディー語 (hi)
• インドネシア語 (id)
• イタリア語 (it)
• 日本語 (ja)
• 韓国語 (ko)
• ペルシャ語 (fa)
• ポルトガル語 (pt)
• ロシア語 (ru)
• スペイン語 (es)
• スワヒリ語 (sw)
• テルグ語 (te)
• タイ語 (th)
• ヨルバ語 (yo)

ライセンス

このモデルは cc-by-nc-4.0 ライセンスの下で提供されています。

パイプラインタグ

sentence-similarity

ライブラリ名

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 が使用されます。これには2つの特徴があります。

• 一方向の注意機構ではなく、双方向の注意機構を使用します。
• クエリテキストに "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: " を使用してプロンプト文字列を手動で指定する必要があります。

DRAMAモデルは、マトリョーシカ表現学習 (MRL) を使用して訓練されており、柔軟な次元数をサポートしています。クエリとドキュメントの両方を、以下のように256などの小さい次元数にエンコードすることができます。

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 の行に対応しています。

📄 ライセンス

このモデルは 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}
}