updating model.py

model.py

@@ -1,15 +1,36 @@
-import torch
-import torch.nn as nn
-
-class BiLSTMClassifier(nn.Module):
-    def __init__(self, vocab_size, embed_dim, hidden_dim, num_classes):
-        super().__init__()
-        self.embedding = nn.Embedding(vocab_size, embed_dim)
-        self.lstm = nn.LSTM(embed_dim, hidden_dim, batch_first=True, bidirectional=True)
-        self.fc = nn.Linear(hidden_dim * 2, num_classes)
-
-    def forward(self, x):
-        x = self.embedding(x)
-        _, (h_n, _) = self.lstm(x)
-        h_cat = torch.cat((h_n[0], h_n[1]), dim=1)
-        return self.fc(h_cat)
+from transformers import PreTrainedModel, PretrainedConfig
+import torch.nn as nn
+import torch
+import torch.nn.functional as F
+
+class BiLSTMConfig(PretrainedConfig):
+    model_type = "bilstm"
+
+    def __init__(self, vocab_size=64000, embedding_dim=1024, hidden_dim=512, num_labels=3, **kwargs):
+        super().__init__(**kwargs)
+        self.vocab_size = vocab_size
+        self.embedding_dim = embedding_dim
+        self.hidden_dim = hidden_dim
+        self.num_labels = num_labels
+
+class BiLSTMClassifier(PreTrainedModel):
+    config_class = BiLSTMConfig
+
+    def __init__(self, config: BiLSTMConfig):
+        super().__init__(config)
+        self.embedding = nn.Embedding(config.vocab_size, config.embedding_dim)
+        self.lstm = nn.LSTM(config.embedding_dim, config.hidden_dim, batch_first=True, bidirectional=True)
+        self.fc = nn.Linear(config.hidden_dim * 2, config.num_labels)
+
+        self.post_init()  
+
+    def forward(self, input_ids, attention_mask=None, labels=None):
+        x = self.embedding(input_ids)
+        _, (h_n, _) = self.lstm(x)
+        h_cat = torch.cat((h_n[0], h_n[1]), dim=1)
+        logits = self.fc(h_cat)
+
+        if labels is not None:
+            loss = F.cross_entropy(logits, labels)
+            return {"loss": loss, "logits": logits}
+        return {"logits": logits}