Upload event_detection_model.py

event_detection_model.py

@@ -0,0 +1,29 @@
+import torch
+from transformers import AutoModel
+from transformers import AutoModelForMaskedLM
+
+class DistillBERTClass(torch.nn.Module):
+    def __init__(self, checkpoint_model):
+        #the super class is not important here!
+        super(DistillBERTClass, self).__init__()
+        #check the rmodel used here !
+        self.pre_trained_model = AutoModelForMaskedLM.from_pretrained(checkpoint_model,output_hidden_states=True)
+        self.linear = torch.nn.Linear(768, 768)
+        self.relu = torch.nn.ReLU()
+        self.dropout = torch.nn.Dropout(0.3)
+        self.classifier = torch.nn.Linear(768, 12)
+
+    def forward(self, input_ids, attention_mask):
+        pre_trained_output = self.pre_trained_model(input_ids=input_ids, attention_mask=attention_mask)
+        hidden_state = pre_trained_output.hidden_states[-1]
+
+        hidden_state = hidden_state[:, 0, :]
+        output = self.linear(hidden_state)
+        output = self.relu(output)
+        output = self.dropout(output)
+        output = self.classifier(output)
+        return output
+
+
+
+