|
import torch
|
|
import torch.nn as nn
|
|
from transformers import AutoModel
|
|
|
|
|
|
class Model(torch.nn.Module):
|
|
|
|
def __init__(self, model_dir, dropout=0.2, hidden_dim=768):
|
|
"""
|
|
Initialize the model.
|
|
:param model_name: the name of the model
|
|
:param metric_names: the names of the metrics to use
|
|
:param dropout: the dropout rate
|
|
:param hidden_dim: the hidden dimension of the model
|
|
"""
|
|
super(Model, self).__init__()
|
|
self.metric_names = ['Happiness', 'Sadness', 'Anger', 'Disgust', 'Fear', 'Pride', 'Valence', 'Arousal']
|
|
self.bert = AutoModel.from_pretrained(model_dir)
|
|
|
|
|
|
for name in self.metric_names:
|
|
setattr(self, name, nn.Linear(hidden_dim, 1))
|
|
setattr(self, 'l_1_' + name, nn.Linear(hidden_dim, hidden_dim))
|
|
|
|
self.layer_norm = nn.LayerNorm(hidden_dim)
|
|
self.relu = nn.ReLU()
|
|
self.dropout = nn.Dropout(dropout)
|
|
self.sigmoid = nn.Sigmoid()
|
|
|
|
def forward(self, input_id, mask):
|
|
"""
|
|
Forward pass of the model.
|
|
:param args: the inputs
|
|
:return: the outputs
|
|
"""
|
|
_, x = self.bert(input_ids = input_id, attention_mask=mask, return_dict=False)
|
|
output = self.rate_embedding(x)
|
|
return output
|
|
|
|
def rate_embedding(self, x):
|
|
output_ratings = []
|
|
for name in self.metric_names:
|
|
first_layer = self.relu(self.dropout(self.layer_norm(getattr(self, 'l_1_' + name)(x) + x)))
|
|
second_layer = self.sigmoid(getattr(self, name)(first_layer))
|
|
output_ratings.append(second_layer)
|
|
|
|
return output_ratings
|
|
|
|
def save_pretrained(self, save_directory):
|
|
self.bert.save_pretrained(save_directory)
|
|
torch.save(self.state_dict(), f'{save_directory}/pytorch_model.bin')
|
|
|
|
@classmethod
|
|
def from_pretrained(cls, model_dir, dropout=0.2, hidden_dim=768):
|
|
model = cls(model_dir, dropout, hidden_dim)
|
|
state_dict = torch.load(f'{model_dir}/pytorch_model.bin', map_location=torch.device('cpu'))
|
|
model.load_state_dict(state_dict)
|
|
return model
|
|
|
