| import torch | |
| from torch import nn | |
| from transformers import MobileBert | |
| class SimModel(PreTrainedModel): | |
| def __init__(self, config): | |
| super().__init__(config) | |
| self.config = config | |
| self.encoder = MobileBert(config) | |
| # Initialize weights and apply final processing | |
| self.post_init() | |
| def forward(self, **input_args): | |
| return self.t5_encoder(**input_args).last_hidden_state |