| import torch | |
| from torch import nn | |
| from transformers import PreTrainedModel, 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) | |
| return self.encoder(input_ids, attention_mask).last_hidden_state |