Update model_script.py

model_script.py

@@ -1,58 +1,70 @@
-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
+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