| import torch |
| import torch.nn as nn |
| from transformers import PreTrainedModel, AutoModel |
| from .model_config import CustomConfig |
|
|
| class LogRegClassifier(nn.Module): |
| def __init__(self, transformer_output_dim): |
| super(LogRegClassifier, self).__init__() |
| self.linear = nn.Linear(transformer_output_dim, 1) |
|
|
| def forward(self, x): |
| return torch.sigmoid(self.linear(x)) |
|
|
| class CombinedModel(PreTrainedModel): |
| config_class = CustomConfig |
|
|
| def __init__(self, config): |
| super().__init__(config) |
| self.transformer = AutoModel.from_pretrained(config.transformer_type) |
| self.classifier = LogRegClassifier(config.transformer_output_dim) |
|
|
| def forward(self, input_ids, attention_mask): |
| outputs = self.transformer(input_ids=input_ids, attention_mask=attention_mask) |
| pooled_output = outputs.last_hidden_state[:, 0, :] |
| return self.classifier(pooled_output) |
|
|
