|
import torch |
|
import torch.nn as nn |
|
import os |
|
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): |
|
if not os.path.isdir(model_dir): |
|
raise ValueError(f"The provided model directory {model_dir} is not a valid directory.") |
|
|
|
model_path = os.path.join(model_dir, 'pytorch_model.bin') |
|
if not os.path.isfile(model_path): |
|
raise FileNotFoundError(f"The model file pytorch_model.bin was not found in the directory {model_dir}.") |
|
|
|
config_path = os.path.join(model_dir, 'config.json') |
|
if not os.path.isfile(config_path): |
|
raise FileNotFoundError(f"The configuration file config.json was not found in the directory {model_dir}.") |
|
|
|
model = cls(model_dir, dropout, hidden_dim) |
|
state_dict = torch.load(model_path, map_location=torch.device('cpu')) |
|
model.load_state_dict(state_dict) |
|
return model |
