No Simple way to add a ValueHead on top of existing HuggingFace Model while Preserving all PreTrainedModel Functionalities?
I have a value network consisting of a LLama2-7b as the base model. I want to attach a Linear Value Head on top. I know I can create a newnn.Module. But, I don’t like it as this way I will lose all nice features of the HuggingFace model likeself.save_pretrained, etc. Is there a way to just add a layer or two while preserving all those properties?I know I may be a bit lost as I am pretty new to this. So, sorry if this sounds like an obvious task. But, I looked at for exampleTRL-Value Headand I see they are doing some complex stuff to achieve this while it seems like an easy task.  Can’t figure out what is prohibiting a simple change.