| import torch | |
| import torch.nn as nn | |
| from transformers import RobertaTokenizer, RobertaModel | |
| from transformers.modeling_outputs import SequenceClassifierOutput | |
| class TargetSentimentClassifier(nn.Module): | |
| def __init__(self, model_name, num_labels=3, dropout_rate=0.3, num_dropouts=5, use_multi_sample_dropout=False, device='cpu'): | |
| super().__init__() | |
| self.device = device | |
| self.tokenizer = RobertaTokenizer.from_pretrained(model_name) | |
| self.roberta = RobertaModel.from_pretrained(model_name) | |
| self.use_multi_sample_dropout = use_multi_sample_dropout | |
| self.num_labels = num_labels | |
| if self.use_multi_sample_dropout: | |
| self.dropouts = nn.ModuleList([nn.Dropout(dropout_rate) for _ in range(num_dropouts)]) | |
| else: | |
| self.dropout = nn.Dropout(dropout_rate) | |
| self.classifier = nn.Linear(self.roberta.config.hidden_size, num_labels) | |
| self.to(device) | |
| def forward(self, input_ids, attention_mask=None, labels=None): | |
| outputs = self.roberta(input_ids, attention_mask=attention_mask) | |
| sequence_output = outputs.last_hidden_state | |
| cls_output = sequence_output[:, 0, :] | |
| if self.use_multi_sample_dropout: | |
| logits_list = [] | |
| for dropout in self.dropouts: | |
| dropped = dropout(cls_output) | |
| logits_list.append(self.classifier(dropped)) | |
| avg_logits = torch.mean(torch.stack(logits_list), dim=0) | |
| else: | |
| dropped = self.dropout(cls_output) | |
| avg_logits = self.classifier(dropped) | |
| loss = None | |
| if labels is not None: | |
| loss_fct = nn.CrossEntropyLoss() | |
| loss = loss_fct(avg_logits.view(-1, self.num_labels), labels.view(-1)) | |
| return SequenceClassifierOutput(loss=loss, logits=avg_logits) | |
| def prepare_input(self, sentence, target, entity_type): | |
| return sentence.replace(target, f"<en> {target} <|{entity_type}|> </en>") | |
| def predict(self, sentence, target, entity_type): | |
| self.eval() | |
| input_text = self.prepare_input(sentence, target, entity_type) | |
| inputs = self.tokenizer( | |
| input_text, | |
| padding=True, | |
| truncation=True, | |
| max_length=512, | |
| return_tensors="pt" | |
| ).to(self.device) | |
| with torch.no_grad(): | |
| outputs = self.forward(**inputs) | |
| logits = outputs.logits | |
| pred = torch.argmax(logits, dim=1).item() | |
| return pred | |