Minor changes for correct inference (#1)

- Update AutoModel... in config.json (dd59cdc4b01d3a8ad2681ed0013ceebc31725bc5)
- Add base_model_prefix="model" on PreTrainedModel (12e66dc10a87bfe0625eccf988d7e1a4f243bbdf)
- Cast attention_mask to bool in SDPA (30e2384fbf7e72648008059a67882cee811cbfd8)
- Tag this model as transformers-compatible (364b6518d26f4b5a6445cb7e350e49427a1cb2c0)
- Specify that we don't want the token_type_ids from the tokenizer (05bfcead01ea0a78660206db4699db8f0791386a)

README.md

@@ -4,6 +4,7 @@ datasets:
 - tiiuae/falcon-refinedweb
 language:
 - en
+library_name: transformers
 ---
 
 # NeoBERT

config.json

@@ -4,7 +4,9 @@
   ],
   "auto_map": {
     "AutoConfig": "model.NeoBERTConfig",
-    "AutoModel": "model.NeoBERTLMHead"
+    "AutoModel": "model.NeoBERT",
+    "AutoModelForMaskedLM": "model.NeoBERTLMHead",
+    "AutoModelForSequenceClassification": "model.NeoBERTForSequenceClassification"
   },
   "classifier_init_range": 0.02,
   "decoder_init_range": 0.02,

model.py

@@ -190,7 +190,7 @@ class EncoderBlock(nn.Module):
                 query=xq.transpose(1, 2),
                 key=xk.transpose(1, 2),
                 value=xv.transpose(1, 2),
-                attn_mask=attention_mask,
+                attn_mask=attention_mask.bool(),
                 dropout_p=0,
             ).transpose(1, 2)
 
@@ -199,6 +199,7 @@ class EncoderBlock(nn.Module):
 
 class NeoBERTPreTrainedModel(PreTrainedModel):
     config_class = NeoBERTConfig
+    base_model_prefix = "model"
     _supports_cache_class = True
 
     def _init_weights(self, module):

tokenizer_config.json

@@ -46,6 +46,10 @@
   "do_lower_case": true,
   "extra_special_tokens": {},
   "mask_token": "[MASK]",
+  "model_input_names": [
+    "input_ids",
+    "attention_mask"
+  ],
   "model_max_length": 4096,
   "pad_token": "[PAD]",
   "sep_token": "[SEP]",