Relaxing `PreTrainedModel` requirement in _save

🚀 Feature request

It’s great to see that Trainer is becoming flexible. Each functions seems to be more self contained now making inheritance easier. I’ve experimented with many custom models. For instance,

class Model(nn.Module):
    def __init__(self, ..):
        self.encoder = AutoModel.from_pretrained(..)
        self.custom_modules = ..
    def forward(self, **kwargs):
        output = self.encoder(**kwargs)
        # some custom operations

Many users are required to create custom models if they just don’t want simple SequenceClassification head. In all cases, I have to override _save method because of this line which explicitly puts a restriction on Trainer to be used with models that inherit from PreTrainedModel. It would be good to relax this requirement and give a warning about not using PreTrainedModel instead.

Your contribution

I’ll open a PR if I get approval.

[sgugger]

After some internal discussion with @julien-c we will lower the requirement from PreTrainedModel to some lower abstractclass/protocol so the user knows exactly what they have to implement for their model to work seamlessly with Trainer. I will work on this end of this week beginning of next.

[prajjwal1]

Sounds good. I’ll look forward to that part then.