Update modeling.py

modeling.py

@@ -28,6 +28,9 @@ class ST2ModelV2(nn.Module):
         
         self.classifier = nn.Linear(self.config.hidden_size, 18)
         #self.classifier = nn.Linear(self.config.hidden_size, 12)
+
+        if self.args.signal_classification and not self.args.pretrained_signal_detector:
+            self.signal_classifier = nn.Linear(self.config.hidden_size, 2)
         
         
 
@@ -73,12 +76,14 @@ class ST2ModelV2(nn.Module):
         #logits = logits.view(batch_size, seq_len, 2, 6).permute(0, 2, 1, 3)
 
         
-        print(logits.shape)
-        st.write(logits.shape)
+        #print(logits.shape)
+        #st.write(logits.shape)
         
         start_arg0_logits, end_arg0_logits, start_arg1_logits, end_arg1_logits, start_sig_logits, end_sig_logits = logits.unbind(-1)
 
-    
+        signal_classification_logits = None
+        if self.args.signal_classification and not self.args.pretrained_signal_detector:
+            signal_classification_logits = self.signal_classifier(sequence_output[:, 0, :])
     
         
         
@@ -90,6 +95,7 @@ class ST2ModelV2(nn.Module):
                 'end_arg1_logits': end_arg1_logits,
                 'start_sig_logits': start_sig_logits,
                 'end_sig_logits': end_sig_logits,
+                'signal_classification_logits': signal_classification_logits,
             }
 
     
@@ -104,8 +110,8 @@ class ST2ModelV2(nn.Module):
         word_ids,
         rel_idx,
     ):
-        print(start_cause_logits.shape)
-        st.write(start_cause_logits.shape)
+        #print(start_cause_logits.shape)
+        #st.write(start_cause_logits.shape)
     
         # Mask special tokens (CLS and SEP)
         for logits in [start_cause_logits, start_effect_logits, end_cause_logits, end_effect_logits]: