Minor changes for correct inference

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]",