| 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) |