I have a (complicated) custom head type which I would like to test with multiple base models. I would like to achieve the following:
- Easily apply this custom head to different base models
- Define custom functions (similar to
generate()) on this custom head - Be able to do
from_pretrained()to load the base model and randomly initialise the new head - Be able to save and load the complete finetuned model with custom head using the
from_pretrained()mechanism
Points 1 and 2 point me towards making a wrapper model.
Points 3 and 4 point me towards inheriting from the base model.
What’s the best practice in this situation? Is there some trick to make the from_pretrained() magic work with the wrapper approach?
One option would be to make a function which automatically creates an inherited model given the base class, but that doesn’t seem like the neat thing to do.