Update modeling_st2.py

modeling_st2.py

@@ -62,16 +62,7 @@ class ST2ModelV2(nn.Module):
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
     ):
-        r"""
-        start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
-            Labels for position (index) of the start of the labelled span for computing the token classification loss.
-            Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
-            are not taken into account for computing the loss.
-        end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
-            Labels for position (index) of the end of the labelled span for computing the token classification loss.
-            Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
-            are not taken into account for computing the loss.
-        """
+        
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
         
         
@@ -116,9 +107,7 @@ class ST2ModelV2(nn.Module):
             signal_classification_logits = self.signal_classifier(sequence_output[:, 0, :])
         # start_logits = start_logits.squeeze(-1).contiguous()
         # end_logits = end_logits.squeeze(-1).contiguous()
-
-        
-                
+      
 
         return {
             'start_arg0_logits': start_arg0_logits,
@@ -131,81 +120,6 @@ class ST2ModelV2(nn.Module):
             
         }
 
-  
-    
-    def position_selector(
-        self,
-        start_cause_logits, 
-        start_effect_logits, 
-        end_cause_logits, 
-        end_effect_logits,
-        attention_mask,
-        word_ids,
-    ):
-        # basic post processing (removing logits from [CLS], [SEP], [PAD])
-        start_cause_logits -= (1 - attention_mask) * 1e4
-        end_cause_logits -= (1 - attention_mask) * 1e4
-        start_effect_logits -= (1 - attention_mask) * 1e4
-        end_effect_logits -= (1 - attention_mask) * 1e4
-
-        start_cause_logits[0] = -1e4
-        end_cause_logits[0] = -1e4
-        start_effect_logits[0] = -1e4
-        end_effect_logits[0] = -1e4
-
-        start_cause_logits[len(word_ids) - 1] = -1e4
-        end_cause_logits[len(word_ids) - 1] = -1e4
-        start_effect_logits[len(word_ids) - 1] = -1e4
-        end_effect_logits[len(word_ids) - 1] = -1e4
-
-        start_cause_logits = torch.log(torch.softmax(start_cause_logits, dim=-1))
-        end_cause_logits = torch.log(torch.softmax(end_cause_logits, dim=-1))
-        start_effect_logits = torch.log(torch.softmax(start_effect_logits, dim=-1))
-        end_effect_logits = torch.log(torch.softmax(end_effect_logits, dim=-1))
-
-        max_arg0_before_arg1 = None
-        for i in range(len(end_cause_logits)):
-            if attention_mask[i] == 0:
-                break
-            for j in range(i + 1, len(start_effect_logits)):
-                if attention_mask[j] == 0:
-                    break
-
-                if max_arg0_before_arg1 is None:
-                    max_arg0_before_arg1 = ((i, j), end_cause_logits[i] + start_effect_logits[j])
-                else:
-                    if end_cause_logits[i] + start_effect_logits[j] > max_arg0_before_arg1[1]:
-                        max_arg0_before_arg1 = ((i, j), end_cause_logits[i] + start_effect_logits[j])
-        
-        max_arg0_after_arg1 = None
-        for i in range(len(end_effect_logits)):
-            if attention_mask[i] == 0:
-                break
-            for j in range(i + 1, len(start_cause_logits)):
-                if attention_mask[j] == 0:
-                    break
-                if max_arg0_after_arg1 is None:
-                    max_arg0_after_arg1 = ((i, j), start_cause_logits[j] + end_effect_logits[i])
-                else:
-                    if start_cause_logits[j] + end_effect_logits[i] > max_arg0_after_arg1[1]:
-                        max_arg0_after_arg1 = ((i, j), start_cause_logits[j] + end_effect_logits[i])
-
-        if max_arg0_before_arg1[1].item() > max_arg0_after_arg1[1].item():
-            end_cause, start_effect = max_arg0_before_arg1[0]
-            start_cause_logits[end_cause + 1:] = -1e4
-            start_cause = start_cause_logits.argmax().item()
-
-            end_effect_logits[:start_effect] = -1e4
-            end_effect = end_effect_logits.argmax().item()
-        else:
-            end_effect, start_cause = max_arg0_after_arg1[0]
-            end_cause_logits[:start_cause] = -1e4
-            end_cause = end_cause_logits.argmax().item()
-
-            start_effect_logits[end_effect + 1:] = -1e4
-            start_effect = start_effect_logits.argmax().item()
-        
-        return start_cause, end_cause, start_effect, end_effect
 
 
     def beam_search_position_selector(
@@ -216,7 +130,6 @@ class ST2ModelV2(nn.Module):
         end_effect_logits,
         topk=5
     ):
-        # basic post processing (removing logits from [CLS], [SEP], [PAD])
 
         start_cause_logits = torch.log(torch.softmax(start_cause_logits, dim=-1))
         end_cause_logits = torch.log(torch.softmax(end_cause_logits, dim=-1))