M2 BERT 8k Retrieval Encoder V1

🚀 Monarch Mixer-BERT

論文Benchmarking and Building Long-Context Retrieval Models with LoCo and M2-BERT のM2-BERT-8k用の80Mチェックポイントです。

このモデルのダウンロードとファインチューニングの方法については、GitHub をご覧ください！

🚀 クイックスタート

💻 使用例

基本的な使用法

このモデルはHugging Faceの AutoModel を使用してロードできます。

from transformers import AutoModelForMaskedLM, BertConfig
config = BertConfig.from_pretrained("hazyresearch/M2-BERT-8K-Retrieval-Encoder-V1")
model = AutoModelForMaskedLM.from_pretrained("hazyresearch/M2-BERT-8K-Retrieval-Encoder-V1", config=config,trust_remote_code=True)

このモデルはHugging Faceの bert-base-uncased tokenizer を使用します。

from transformers import BertTokenizer
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')

高度な使用法

このモデルは検索用の埋め込みを生成します。埋め込みの次元数は768です。

from transformers import AutoTokenizer, AutoModelForMaskedLM, BertConfig

max_seq_length = 8192
testing_string = "Every morning, I make a cup of coffee to start my day."
config = BertConfig.from_pretrained("hazyresearch/M2-BERT-8K-Retrieval-Encoder-V1")
model = AutoModelForMaskedLM.from_pretrained("hazyresearch/M2-BERT-8K-Retrieval-Encoder-V1", config=config, trust_remote_code=True)

tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased", model_max_length=max_seq_length)
input_ids = tokenizer([testing_string], return_tensors="pt", padding="max_length", return_token_type_ids=False, truncation=True, max_length=max_seq_length)

outputs = model(**input_ids)
embeddings = outputs['sentence_embedding']

🔧 技術詳細

リモートコード

このモデルを from_pretrained メソッドでロードする際には、trust_remote_code=True を渡す必要があります。これは、独自のPyTorchコードを使用しているためです（詳細はGitHubを参照）。コードの正確なgitコミットを指定する revision 引数を渡すことを検討してください。

mlm = AutoModelForMaskedLM.from_pretrained(
   "hazyresearch/M2-BERT-8K-Retrieval-Encoder-V1",
   config=config,
   trust_remote_code=True,
)

設定

use_flash_mm はデフォルトで false です。現在、FlashMMの使用はサポートされていません。

📄 ライセンス

このプロジェクトはApache-2.0ライセンスの下で公開されています。