Update modeling_st2.py

modeling_st2.py

@@ -41,23 +41,7 @@ class ST2ModelV2(nn.Module):
         self.dropout = nn.Dropout(classifier_dropout)
         self.classifier = nn.Linear(self.config.hidden_size, 6)
 
-        if args.mlp:
-            self.classifier = nn.Sequential(
-                nn.Linear(self.config.hidden_size, self.config.hidden_size),
-                nn.ReLU(),
-                nn.Linear(self.config.hidden_size, 6),
-                nn.Tanh(),
-                nn.Linear(6, 6),
-            )
-
-        if args.add_signal_bias:
-            self.signal_phrases_layer = nn.Parameter(
-                torch.normal(
-                    mean=self.model.embeddings.word_embeddings.weight.data.mean(), 
-                    std=self.model.embeddings.word_embeddings.weight.data.std(),
-                    size=(1, self.config.hidden_size),
-                )
-            )
+        
         
         if self.args.signal_classification and not self.args.pretrained_signal_detector:
             self.signal_classifier = nn.Linear(self.config.hidden_size, 2)
@@ -89,37 +73,22 @@ class ST2ModelV2(nn.Module):
         """
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
         
-        if signal_bias_mask is not None and not self.args.signal_bias_on_top_of_lm:
-            inputs_embeds = self.signal_phrases_bias(input_ids, signal_bias_mask)
-
-            outputs = self.model(
-                # input_ids,
-                attention_mask=attention_mask,
-                token_type_ids=token_type_ids,
-                position_ids=position_ids,
-                head_mask=head_mask,
-                inputs_embeds=inputs_embeds,
-                output_attentions=output_attentions,
-                output_hidden_states=output_hidden_states,
-                return_dict=return_dict,
-            )
-        else:
+        
             
-            outputs = self.model(
-                input_ids,
-                attention_mask=attention_mask,
-                token_type_ids=token_type_ids,
-                position_ids=position_ids,
-                head_mask=head_mask,
-                inputs_embeds=inputs_embeds,
-                output_attentions=output_attentions,
-                output_hidden_states=output_hidden_states,
-                return_dict=return_dict,
-            )
+        outputs = self.model(
+            input_ids,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
 
         sequence_output = outputs[0]
-        if signal_bias_mask is not None and self.args.signal_bias_on_top_of_lm:
-            sequence_output[signal_bias_mask == 1] += self.signal_phrases_layer
+        
 
         sequence_output = self.dropout(sequence_output)
         logits = self.classifier(sequence_output)  # [batch_size, max_seq_length, 6]
@@ -147,36 +116,7 @@ class ST2ModelV2(nn.Module):
         # start_logits = start_logits.squeeze(-1).contiguous()
         # end_logits = end_logits.squeeze(-1).contiguous()
 
-        arg0_loss = None
-        arg1_loss = None
-        sig_loss = None
-        total_loss = None
-        signal_classification_loss = None
-        if start_positions is not None and end_positions is not None:
-            loss_fct = nn.CrossEntropyLoss()
-
-            start_arg0_loss = loss_fct(start_arg0_logits, start_positions[:, 0])
-            end_arg0_loss = loss_fct(end_arg0_logits, end_positions[:, 0])
-            arg0_loss = (start_arg0_loss + end_arg0_loss) / 2
-
-            start_arg1_loss = loss_fct(start_arg1_logits, start_positions[:, 1])
-            end_arg1_loss = loss_fct(end_arg1_logits, end_positions[:, 1])
-            arg1_loss = (start_arg1_loss + end_arg1_loss) / 2
-
-            # sig_loss = 0.
-            start_sig_loss = loss_fct(start_sig_logits, start_positions[:, 2])
-            end_sig_loss = loss_fct(end_sig_logits, end_positions[:, 2])
-            sig_loss = (start_sig_loss + end_sig_loss) / 2
-
-            if sig_loss.isnan():
-                sig_loss = 0.
-
-            if self.args.signal_classification and not self.args.pretrained_signal_detector:
-                signal_classification_labels = end_positions[:, 2] != -100
-                signal_classification_loss = loss_fct(signal_classification_logits, signal_classification_labels.long())
-                total_loss = (arg0_loss + arg1_loss + sig_loss + signal_classification_loss) / 4
-            else:
-                total_loss = (arg0_loss + arg1_loss + sig_loss) / 3
+        
                 
 
         return {
@@ -187,69 +127,10 @@ class ST2ModelV2(nn.Module):
             'start_sig_logits': start_sig_logits,
             'end_sig_logits': end_sig_logits,
             'signal_classification_logits': signal_classification_logits,
-            'arg0_loss': arg0_loss,
-            'arg1_loss': arg1_loss,
-            'sig_loss': sig_loss,
-            'signal_classification_loss': signal_classification_loss,
-            'loss': total_loss,
+            
         }
 
-    """
-    def save_pretrained(self, save_directory):
-        
-        #Save model state dict as safetensor, configuration, and tokenizer files.
-        
-        # Ensure the directory exists
-        os.makedirs(save_directory, exist_ok=True)
-
-        # Save model state dict as safetensor (use torch.save for PyTorch model)
-        model_path = os.path.join(save_directory, "model.safetensor")
-        save_file(self.state_dict(), model_path)
-
-        # Save config if available
-        config_save_path = os.path.join(save_directory, 'config.json')
-        self.config.to_json_file(config_save_path)
-        
-
-        # Save tokenizer 
-        if hasattr(self, 'tokenizer') and self.tokenizer is not None:
-            tokenizer_save_path = os.path.join(save_directory, 'tokenizer')
-            self.tokenizer.save_pretrained(tokenizer_save_path)
-    """
-    
-
-    def save_pretrained(self, save_directory):
-        """
-        Save model state dict as safetensor, PyTorch .bin format, configuration, and tokenizer files.
-        """
-        # Ensure the directory exists
-        os.makedirs(save_directory, exist_ok=True)
-    
-        # Save model state dict as safetensor
-        model_path_safetensor = os.path.join(save_directory, "model.safetensors")
-        save_file(self.state_dict(), model_path_safetensor)  # Save as .safetensors
-    
-        # Save model state dict as PyTorch .bin (traditional format)
-        model_path_bin = os.path.join(save_directory, "pytorch_model.bin")
-        torch.save(self.state_dict(), model_path_bin)  # Save as .bin using PyTorch's torch.save()
-    
-        # Save config if available
-        config_save_path = os.path.join(save_directory, 'config.json')
-        self.config.to_json_file(config_save_path)
-
-        """
-        # Save tokenizer if it exists
-        if hasattr(self, 'tokenizer') and self.tokenizer is not None:
-            tokenizer_save_path = os.path.join(save_directory, 'tokenizer')
-            self.tokenizer.save_pretrained(tokenizer_save_path)
-        """
-
-
-    def signal_phrases_bias(self, input_ids, signal_bias_mask):
-        inputs_embeds = self.model.get_input_embeddings()(input_ids)
-        inputs_embeds[signal_bias_mask == 1] += self.signal_phrases_layer  # self.signal_phrases_layer(inputs_embeds[signal_bias_mask == 1])
-
-        return inputs_embeds
+  
     
     def position_selector(
         self,
@@ -332,27 +213,11 @@ class ST2ModelV2(nn.Module):
         start_effect_logits, 
         end_cause_logits, 
         end_effect_logits,
-        attention_mask,
         word_ids,
         topk=5
     ):
         # 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))
@@ -360,19 +225,12 @@ class ST2ModelV2(nn.Module):
 
         scores = dict()
         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
                 scores[str((i, j, "before"))] = end_cause_logits[i].item() + start_effect_logits[j].item()
         
         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
                 scores[str((i, j, "after"))] = start_cause_logits[j].item() + end_effect_logits[i].item()