chandar-lab

/

NeoBERT

like 57

Follow

Chandar Research Lab 20

Hello @Lolalb and team!

Preface

Congratulations on your model release! Very exciting to see another strong encoder out there. I already started some training runs of my own to experiment! I ran into some minor issues regarding the transformers integration, that this PR tackles.

Pull Request overview

• Update AutoModel... in config.json
• Add base_model_prefix="model" on PreTrainedModel
• Cast attention_mask to bool in SDPA
• Tag this model as transformers-compatible
• Specify that we don't want the token_type_ids from the tokenizer

Details

Updating config.json

Previously, loading an AutoModel.from_pretrained actually loaded the Masked language modeling model, and it was not possible to load the Sequence Classification model. Together with updating the base_model_prefix in model.py, it's now possible to do AutoModel.from_pretrained, AutoModelForMaskedLM.from_pretrained and AutoModelForSequenceClassification.from_pretrained.

Cast attention mask to bool

In SDPA, the attention mask was an int tensor before, which SDPA doesn't like. I'm doing some training with this model now, but I have a custom data collator, so it's not using your data collator for Flash Attention 2 support.

No token_type_ids from the tokenizer

The BERT tokenizer outputs token_type_ids by default. This used to be used to recognize which sentence was which if you provided pairs of sentences. However, it's kind of fallen out of style. We don't need it, but the tokenizer does return it, so this removes that.

You can test all this stuff nicely by using the revision argument:

from transformers import AutoModel, AutoTokenizer

model_name = "chandar-lab/NeoBERT"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True, revision="refs/pr/1")
model = AutoModel.from_pretrained(model_name, trust_remote_code=True, revision="refs/pr/1")

# Tokenize input text
text = "NeoBERT is the most efficient model of its kind!"
inputs = tokenizer(text, return_tensors="pt")

# Generate embeddings
outputs = model(**inputs)
embedding = outputs.last_hidden_state[:, 0, :]
print(embedding.shape)

I'll be uploading my first finetune of this model in a few minutes. It looks to be very strong, stronger than my same model with ModernBERT-base.
Edit: Here's the NeoBERT model and my ModernBERT-base baseline, using the same training script:

• https://huggingface.co/tomaarsen/reranker-ModernBERT-base-gooaq-bce-static-retriever-hardest
• https://huggingface.co/tomaarsen/reranker-NeoBERT-gooaq-bce

Interestingly, NeoBERT is much stronger in-domain, but much worse out-of-domain with my finetuned model here.

If there is a lot of interest from the community, it might make sense to introduce neobert as an architecture in transformers, so that users won't have to use trust_remote_code anymore. I do have to preface that we don't have xformers in pure transformers, so we would need a "manual" swiglu instead.

cc @stefan-it as you were also working on NeoBERT, I believe.

• Tom Aarsen