Tom Aarsen commited on
Commit
dd59cdc
·
1 Parent(s): c10259c

Update AutoModel... in config.json

Browse files

Previously, loading an AutoModel.from_pretrained actually loaded the Masked language modeling model, and it was not possible to load the Sequence Classification model.

Files changed (1) hide show
  1. config.json +3 -1
config.json CHANGED
@@ -4,7 +4,9 @@
4
  ],
5
  "auto_map": {
6
  "AutoConfig": "model.NeoBERTConfig",
7
- "AutoModel": "model.NeoBERTLMHead"
 
 
8
  },
9
  "classifier_init_range": 0.02,
10
  "decoder_init_range": 0.02,
 
4
  ],
5
  "auto_map": {
6
  "AutoConfig": "model.NeoBERTConfig",
7
+ "AutoModel": "model.NeoBERT",
8
+ "AutoModelForMaskedLM": "model.NeoBERTLMHead",
9
+ "AutoModelForSequenceClassification": "model.NeoBERTForSequenceClassification"
10
  },
11
  "classifier_init_range": 0.02,
12
  "decoder_init_range": 0.02,