Extending a pretrained model

Hello, I would like to add some tweaks to DistilBertForQuestionAnswering. I created a new class extended from DistilBertForQuestionAnswering and added few minor tweaks to the classification layer. However I am not sure how to load the pre-trained distilbert-base-uncased model using my extended class. Following code is throwing error:

model = My_extended_DistilBertForQuestionAnswering.from_pretrained("distilbert-base-uncased")

AttributeError: module 'My_extended_DistilBertForQuestionAnswering' has no attribute 'from_pretrained'

Code snippet of extended class:

class My_extended_DistilBertForQuestionAnswering(DistilBertForQuestionAnswering):
    def __init__(self, config):
        super().__init__(config)
        # .... define new linear layer, dropout etc...

This code snippet works fine for me:

from transformers import DistilBertForQuestionAnswering

class My_extended_DistilBertForQuestionAnswering(DistilBertForQuestionAnswering):
    def __init__(self, config):
        super().__init__(config)
        # .... define new linear layer, dropout etc...

model = My_extended_DistilBertForQuestionAnswering.from_pretrained("distilbert-base-uncased")

The fact that your error says the “module” doesn’t have the attribute makes me think you’re importing your custom class incorrectly. Maybe double check that you’re importing your custom class correctly:

from where_the_class_is_defined import My_extended_DistilBertForQuestionAnswering

model = My_extended_DistilBertForQuestionAnswering.from_pretrained("distilbert-base-uncased"
2 Likes