Loading and save different models types in 1 class

🤗Transformers
1

Hello,

I created my custom class:

class LingMessCoref(BertPreTrainedModel):
    def __init__(self, config, args):
        super().__init__(config)

        base_model = AutoModel.from_config(config)
        self.base_model_prefix = base_model.base_model_prefix
        self.config_class = base_model.config_class
        setattr(self, self.base_model_prefix, base_model)

        self.clf = Linear(config.hidden_size, 1)

        self.init_weights()

I did the base_model thing to be able to change the transformer to one of these options: ['longformer', 'roberta', 'deberta', 'bert']

the first initialization is:

model, loading_info = LingMessCoref.from_pretrained(
        args.model_name_or_path, output_loading_info=True,
        config=config, cache_dir=args.cache_dir, args=args
    )

where model_name_or_path = 'allenai/longformer-large-4096'

To save the model checkpoint I run:
model.save_pretrained(output_dir)

Then if I want to load the model again using from_pretrained it seems like the weights not loaded.
the model accuracy and loss is different from the checkpoint run.

Is it make sense that when I change self to LingMessCoref it works ?

base_model = AutoModel.from_config(config)
LingMessCoref.base_model_prefix = base_model.base_model_prefix
LingMessCoref.config_class = base_model.config_class
setattr(self, self.base_model_prefix, base_model)

Any suggestions?

Thanks in advance,
Shon