license: apache-2.0 language:

• en pipeline_tag: fill-mask inference: false

Monarch Mixer-BERT模型

这是论文《Benchmarking and Building Long-Context Retrieval Models with LoCo and M2-BERT》中M2-BERT-128模型的8000万参数检查点。

访问我们的GitHub仓库获取下载和微调指南！

使用方法

可通过Hugging Face的AutoModel加载本模型：

from transformers import AutoModelForMaskedLM
model = AutoModelForMaskedLM.from_pretrained("jonsaadfalcon/M2-BERT-128-Retrieval-Encoder-V1", trust_remote_code=True)

本模型使用Hugging Face的bert-base-uncased分词器：

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

功能应用

本模型用于生成768维的检索嵌入向量：

from transformers import AutoTokenizer, AutoModelForMaskedLM

max_seq_length = 128
测试文本 = "每天早晨，我都会冲杯咖啡开始新的一天。"
model = AutoModelForMaskedLM.from_pretrained("jonsaadfalcon/M2-BERT-128-Retrieval-Encoder-V1", trust_remote_code=True)

tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased", model_max_length=max_seq_length)
input_ids = tokenizer([测试文本], 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']

远程代码说明

加载本模型需设置trust_remote_code=True参数，因为我们使用了自定义PyTorch代码（详见GitHub）。建议通过revision参数指定代码的git提交哈希，例如：

mlm = AutoModelForMaskedLM.from_pretrained(
   "jonsaadfalcon/M2-BERT-128-Retrieval-Encoder-V1",
   trust_remote_code=True,
)

配置说明

注意use_flash_mm参数默认关闭，当前版本暂不支持FlashMM功能。