Saving underlying language model after trained on downstream task

Hello everyone,

Although I am quite new here (first post!), I feel like I know huggingface transformers pretty well since I read the documentation and tutorials, and I even made a small project by myself recently.

My goal is to create a module that receives a pretrained language model, fine tunes/trains it on some downstream task, NER for example, and then saves only the language model but without the NER classifier at the end (only the underlying finetuned language model) .

Some requirements that make this task non-trivial:

  • The module should be general enough (inherit from PretrainedModel) so it can be trained, saved to disk, loaded from disk, and then trained, all using huggingface Trainer.
  • The module must be flexible enough to do this for all huggingface’s models with as little code changes as possible.
  • The module should be as forward compatible as possible, meaning that I prefer using AutoModel as much as I can.

I hope I got myself clear. if you have any questions regarding my post, please don’t hesitate to reply.