| import torch | |
| from torch import nn | |
| from transformers import PreTrainedModel, MobileBertModel | |
| class SimModel(PreTrainedModel): | |
| 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_args): | |
| print(**input_args) | |
| return self.encoder(**input_args).last_hidden_state |