|
|
|
import torch |
|
from transformers import AutoModelForSequenceClassification, AutoTokenizer |
|
|
|
class SentimentAnalysisModel(torch.nn.Module): |
|
def __init__(self, bert_model_name, num_labels=7, dropout_rate=0.4): |
|
super(SentimentAnalysisModel, self).__init__() |
|
self.model = AutoModelForSequenceClassification.from_pretrained(bert_model_name) |
|
self.model.classifier.out_proj = torch.nn.Linear(self.model.config.hidden_size, num_labels) |
|
self.model.classifier.dropout = torch.nn.Dropout(p=dropout_rate, inplace=False) |
|
|
|
for layer in self.model.roberta.encoder.layer: |
|
layer.attention.self.dropout = torch.nn.Dropout(p=dropout_rate) |
|
layer.attention.output.dropout = torch.nn.Dropout(p=dropout_rate) |
|
layer.intermediate.dropout = torch.nn.Dropout(p=dropout_rate) |
|
layer.output.dropout = torch.nn.Dropout(p=dropout_rate) |
|
|
|
self.tokenizer = AutoTokenizer.from_pretrained(bert_model_name) |
|
|
|
def forward(self, input_ids, attention_mask, labels=None): |
|
outputs = self.model(input_ids, attention_mask=attention_mask, labels=labels) |
|
return outputs |
|
|
