README.md

---
license: apache-2.0
language:
- en
pipeline_tag: fill-mask
inference: false
---

# Monarch Mixer-BERT

The 80M checkpoint for M2-BERT-2k from the paper [Benchmarking and Building Long-Context Retrieval Models with LoCo and M2-BERT](https://arxiv.org/abs/2402.07440).

Check out our [GitHub](https://github.com/HazyResearch/m2/tree/main) for instructions on how to download and fine-tune it!

## How to use

You can load this model using Hugging Face `AutoModel`:
```python
from transformers import AutoModelForMaskedLM, BertConfig
config = BertConfig.from_pretrained("hazyresearch/M2-BERT-2K-Retrieval-Encoder-V1")
model = AutoModelForMaskedLM.from_pretrained("hazyresearch/M2-BERT-2k-Retrieval-Encoder-V1", config=config, trust_remote_code=True)
```

This model uses the Hugging Face `bert-base-uncased tokenizer`:
```
from transformers import BertTokenizer
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
```

## How to use

This model generates embeddings for retrieval. The embeddings have a dimensionality of 768:
```
from transformers import AutoTokenizer, AutoModelForMaskedLM, BertConfig

max_seq_length = 2048
testing_string = "Every morning, I make a cup of coffee to start my day."
config = BertConfig.from_pretrained("hazyresearch/M2-BERT-2K-Retrieval-Encoder-V1")
model = AutoModelForMaskedLM.from_pretrained("hazyresearch/M2-BERT-2k-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']
```

### Remote Code

This model requires `trust_remote_code=True` to be passed to the `from_pretrained` method. This is because we use custom PyTorch code (see our GitHub). You should consider passing a `revision` argument that specifies the exact git commit of the code, for example:

```python
mlm = AutoModelForMaskedLM.from_pretrained(
   "hazyresearch/M2-BERT-2k-Retrieval-Encoder-V1",
   config=config,
   trust_remote_code=True,
)
```

### Configuration
Note `use_flash_mm` is false by default. Using FlashMM is currently not supported.