| from transformers import PreTrainedModel, AutoModel | |
| from .configuration_moral import BertItalianoConfig | |
| import torch | |
| class BertItaliano(PreTrainedModel): | |
| config_class = BertItalianoConfig | |
| def __init__(self, config, num_labels_1=6, num_labels_2=3): | |
| super(BertItaliano, self).__init__(config) | |
| self.num_labels1 = num_labels_1 | |
| self.num_labels2 = num_labels_2 | |
| self.bert = AutoModel.from_pretrained("dbmdz/bert-base-italian-xxl-uncased") | |
| self.dropout = torch.nn.Dropout(0.3) | |
| self.linear_1 = torch.nn.Linear(config.hidden_size, num_labels_1) | |
| self.linear_2 = torch.nn.Linear(config.hidden_size, num_labels_2) | |
| def forward(self, input_ids, attention_mask=None, token_type_ids=None): | |
| outputs = self.bert(input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids) | |
| pooled_output = outputs[1] | |
| pooled_output = self.dropout(pooled_output) | |
| logits_1 = self.linear_1(pooled_output) | |
| logits_2 = self.linear_2(pooled_output) | |
| return logits_1, logits_2 |