| import torch |
|
|
| from typing import Tuple |
| from transformers import PreTrainedModel, BertModel |
| from torch import nn |
|
|
|
|
| class BertClassificationModel(PreTrainedModel): |
| def __init__(self, config, num_main_segment=None, num_sub_segment=None): |
| super(BertClassificationModel, self).__init__(config=config) |
| self.num_main_segment = num_main_segment if num_main_segment else config.num_main_segment |
| self.num_sub_segment = num_sub_segment if num_sub_segment else config.num_sub_segment |
| self.bert = BertModel.from_pretrained("bert-base-multilingual-uncased") |
| self.dropout = nn.Dropout(0.1) |
| self.hidden_2 = nn.Linear(768, 768) |
| self.fc_main = nn.Linear(768, self.num_main_segment) |
| self.fc_sub = nn.Linear(768, self.num_sub_segment) |
|
|
| def forward(self, input_ids: torch.tensor, attention_masks: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: |
| outputs = self.bert(input_ids=input_ids, |
| attention_mask=attention_masks) |
| last_hidden_state_cls = outputs[0][:, 0, :] |
| return self.fc_main(last_hidden_state_cls), self.fc_sub(last_hidden_state_cls) |
