Drama 1b

🚀 DRAMA-1B: 大規模言語モデルから小規模な密ベクトル検索器への多様な拡張

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

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

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

🚀 クイックスタート

✨ 主な機能

• 剪定された大規模言語モデルをバックボーンとする密ベクトル検索モデル
• 高品質なデータ拡張により、コンパクトなサイズで英語と多言語検索タスクで強力な性能を発揮
• マトリョーシカ表現学習により、埋め込み次元数を柔軟に調整可能

📦 インストール

このモデルは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-1b"
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.5062, 0.1475], [0.1837, 0.6331]]

⚠️ 重要提示

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.6579, 0.3296], [0.3388, 0.7547]]

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-1b", 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.5062, 0.1475], [0.1837, 0.6331]]

⚠️ 重要提示

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-1b", 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.6579, 0.3296], [0.3388, 0.7547]]

📚 詳細ドキュメント

評価

このモデルは、BEIR、MIRACL、MLDR、およびMTEBのいくつかの多言語検索タスクを含む複数の検索ベンチマークで評価されています。英語と多言語検索タスクの両方で強力な性能を示しています。

このページで公開されているdrama-1bは、10億個の非埋め込みパラメータを持つDRAMA-1Bに対応しています。

サポート言語

DRAMA-1Bは、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}
}