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 new nn.Module. But, I don’t like it as this way I will lose all nice features of the HuggingFace model like self.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 example TRL-Value Head and 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.