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):
|