Tom Aarsen
commited on
Commit
·
12e66dc
1
Parent(s):
dd59cdc
Add base_model_prefix="model" on PreTrainedModel
Browse filesThis allows you to load the AutoModel, even though the model checkpoint is for the Masked Language modeling head. In short: the base_model_prefix is what it searches for in the subclasses like NeoBERTLMHead and NeoBERTForSequenceClassification
model.py
CHANGED
@@ -199,6 +199,7 @@ class EncoderBlock(nn.Module):
|
|
199 |
|
200 |
class NeoBERTPreTrainedModel(PreTrainedModel):
|
201 |
config_class = NeoBERTConfig
|
|
|
202 |
_supports_cache_class = True
|
203 |
|
204 |
def _init_weights(self, module):
|
|
|
199 |
|
200 |
class NeoBERTPreTrainedModel(PreTrainedModel):
|
201 |
config_class = NeoBERTConfig
|
202 |
+
base_model_prefix = "model"
|
203 |
_supports_cache_class = True
|
204 |
|
205 |
def _init_weights(self, module):
|