import torch
from torch import nn
from transformers import PreTrainedMobileBertModel, MobileBertModel

class SimModel(MobileBertModel):
    def __init__(self, config):
        super().__init__(config)
        self.config = config
        self.encoder = MobileBertModel(config)
        # Initialize weights and apply final processing
        self.post_init()

    def forward(self, input_ids, attention_mask, token_type_ids, return_dict):
        print(input_ids, attention_mask, token_type_ids)
        print(return_dict)
        return self.encoder(input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids, return_dict=return_dict)