| import copy | |
| from transformers import PreTrainedModel, PretrainedConfig, AutoModel, AutoTokenizer | |
| import torch.nn as nn | |
| import torch | |
| from typing import List | |
| class AMC_OPT_conf(PretrainedConfig): | |
| model_type = "bert" | |
| def __init__(self, | |
| out_labels=3, | |
| emb_size=768, | |
| drop_out=0.1, | |
| pretrained_model="m-polignano-uniba/bert_uncased_L-12_H-768_A-12_italian_alb3rt0", | |
| **kwargs): | |
| self.out_labels = out_labels | |
| self.drop_out = drop_out | |
| self.emb_size = emb_size | |
| self.pretrained_model = pretrained_model | |
| self.fine_tuned_model = "data/models/amc_opt_msmd_nocoadapt.pt" | |
| super().__init__(**kwargs) | |
| class AMC_OPT_sub(nn.Module): | |
| def __init__(self, config): | |
| super(AMC_OPT_sub, self).__init__() | |
| self.model = copy.deepcopy(AutoModel.from_pretrained(config.pretrained_model)) | |
| self.dropout1 = nn.Dropout(config.drop_out) | |
| self.linear1 = nn.Linear(config.emb_size, config.out_labels) | |
| self.loss_fct = nn.CrossEntropyLoss() | |
| self.hyper_params = {'learning_rate': 6.599917952321265e-05, | |
| 'weight_decay': 0.02157165894420757, | |
| 'warmup_steps': 0.8999999999999999, | |
| 'num_epochs': 11} | |
| self.params_name = "alberto_multiclass_opt_msmd.pt" | |
| def forward(self, labels, input_ids, attention_mask, **args): | |
| outputs = self.model(input_ids=input_ids, attention_mask=attention_mask, **args) | |
| x = self.dropout1(outputs[1]) | |
| logits = self.linear1(x) | |
| if labels != None: | |
| loss = self.loss_fct(logits, labels) | |
| return {'logits':logits, 'loss':loss} | |
| else: | |
| return {'logits':logits} | |
| class AMC_OPT(PreTrainedModel): | |
| config_class = AMC_OPT_conf | |
| def __init__(self, config): | |
| super().__init__(config) | |
| self.model = AMC_OPT_sub(config) | |
| self.model.load_state_dict(torch.load(config.fine_tuned_model)) | |
| def forward(self, labels, input_ids, attention_mask, **args): | |
| return self.model(labels, input_ids, attention_mask, **args) | |
| tokenizer = AutoTokenizer.from_pretrained("m-polignano-uniba/bert_uncased_L-12_H-768_A-12_italian_alb3rt0") | |
| tokenizer.model_max_length = 128 |