[FIXED] How to Integrate Custom OpenAIModel into a AutoModelForSequenceClassification Model?

Issue

I've developed a custom OpenAIModel module that acts like BERT models but makes an OpenAI embeddings request and returns the results when called. I want to use this module, utilizing Hugging Face's transformers library inside an AutoModelForSequenceClassification model. However, when I try to replace the base_model with my OpenAIModel using the following code, the base_model property doesn't change as expected:

from transformers import AutoModelForSequenceClassification
import torch

num_labels = ... # Define the appropriate number of labels
train_arch = ... # Define parameters specific to your training architecture

model = AutoModelForSequenceClassification.from_pretrained('dbmdz/distilbert-base-turkish-cased', num_labels=num_labels)
model.base_model = OpenAIModel(train_arch)
model.classifier = torch.nn.Linear(model.base_model.config.dim, num_labels)
model.config = OpenAIModelConfig()
model.config_class = type(OpenAIModelConfig)
tokenizer = OpenAITokenizer()

After executing model.base_model = OpenAIModel(train_arch), when I inspect model.base_model, it still appears to be distilbert and not the expected OpenAIModel.

I can successfully change all other values of the SequenceClassification model, but the base_model remains unaltered. I'm unsure why this is happening and how to resolve it.

What is the correct way to integrate my model into AutoModelForSequenceClassification? Or should I be using a different approach to change the base_model? (I tried to load my model with from_pretrained but AutoModel class doesn't support custom model structure as i see)

Any help or suggestions would be greatly appreciated!

Solution

If we are using a distilbert model, changing the model.distilbert layer, not model.base_model, was enough to do what I intended.

When I did model.distilbert = OpenAIModel(**kwargs) I achieved my goal.

Answered By - Salihcan