UNCANNY69 commited on
Commit
ce8a9f9
1 Parent(s): 9eb7d2b

Create model.py

Browse files
Files changed (1) hide show
  1. model.py +66 -0
model.py ADDED
@@ -0,0 +1,66 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import PretrainedConfig
2
+ from abc import ABCMeta
3
+ import torch
4
+ from transformers.pytorch_utils import nn
5
+ from transformers import RobertaModel, RobertaConfig
6
+ import torch
7
+ import torch.nn.functional as F
8
+ from transformers import PreTrainedModel
9
+ from transformers.modeling_outputs import SequenceClassifierOutput
10
+ from transformers import PretrainedConfig
11
+
12
+ class RoBERTaLSTMConfig(PretrainedConfig):
13
+ model_type = "robertaLSTMForSequenceClassification"
14
+
15
+ def __init__(self,
16
+ num_classes=2,
17
+ embed_dim=768,
18
+ num_layers=12,
19
+ hidden_dim_lstm=256, # New parameter for LSTM
20
+ dropout_rate=0.1,
21
+ **kwargs):
22
+ super().__init__(**kwargs)
23
+ self.num_classes = num_classes
24
+ self.embed_dim = embed_dim
25
+ self.num_layers = num_layers
26
+ self.hidden_dim_lstm = hidden_dim_lstm # Assign LSTM hidden dimension
27
+ self.dropout_rate = dropout_rate
28
+ self.id2label = {
29
+ 0: "fake",
30
+ 1: "true",
31
+ }
32
+ self.label2id = {
33
+ "fake": 0,
34
+ "true": 1,
35
+ }
36
+
37
+ class RoBERTaLSTMForSequenceClassification(PreTrainedModel, metaclass=ABCMeta):
38
+ def __init__(self, config):
39
+ super(RoBERTaLSTMForSequenceClassification, self).__init__(config)
40
+ self.num_classes = config.num_classes
41
+ self.embed_dim = config.embed_dim
42
+ self.num_layers = config.num_layers
43
+ self.hidden_dim_lstm = config.hidden_dim_lstm
44
+ self.dropout = nn.Dropout(config.dropout_rate)
45
+ self.roberta = RobertaModel.from_pretrained('roberta-base',output_hidden_states=True,
46
+ output_attentions=False)
47
+ print("RoBERTa Model Loaded")
48
+ self.lstm = nn.LSTM(self.embed_dim, self.hidden_dim_lstm, batch_first=True, num_layers=3)
49
+ self.fc = nn.Linear(self.hidden_dim_lstm, self.num_classes)
50
+
51
+ def forward(self, input_ids, attention_mask, labels=None):
52
+ roberta_output = self.roberta(input_ids=input_ids, attention_mask=attention_mask)
53
+ hidden_states = roberta_output.last_hidden_state
54
+ out, _ = self.lstm(hidden_states)
55
+ out = self.dropout(out[:, -1, :])
56
+ logits = self.fc(out)
57
+ loss = None
58
+ if labels is not None:
59
+ loss = F.cross_entropy(logits, labels)
60
+ out = SequenceClassifierOutput(
61
+ loss=loss,
62
+ logits=logits,
63
+ hidden_states=roberta_output.hidden_states,
64
+ attentions=roberta_output.attentions,
65
+ )
66
+ return out