Rename ST2ModelV2_6.py to modeling_st2.py

ST2ModelV2_6.py → modeling_st2.py

@@ -7,79 +7,199 @@ from transformers import (
     AutoConfig, 
     AutoModelForSequenceClassification
     )
-from statistics import mode
-from safetensors.torch import load_file
-
-
+import os
+from safetensors.torch import save_file
+
+
+
+class SignalDetector(nn.Module):
+    def __init__(self, model_and_tokenizer_path) -> None:
+        super().__init__()
+        self.tokenizer = AutoTokenizer.from_pretrained(model_and_tokenizer_path) 
+        self.signal_detector = AutoModelForSequenceClassification.from_pretrained(model_and_tokenizer_path)
+        self.signal_detector.eval()
+        self.signal_detector.cuda()
+    
+    @torch.no_grad()
+    def predict(self, text: str) -> int:
+        input_ids = self.tokenizer.encode(text)
+        input_ids = torch.tensor([input_ids]).cuda()
+        outputs = self.signal_detector(input_ids)
+        return outputs[0].argmax().item()
+
+    
 class ST2ModelV2(nn.Module):
-    def __init__(self, args, config):
+    def __init__(self, args):
         super(ST2ModelV2, self).__init__()
         self.args = args
-        self.config = config
 
-        # Load the base model (e.g., Roberta)
-        self.model = AutoModel.from_pretrained("roberta-large", config=config)
+        self.config = AutoConfig.from_pretrained(args.model_name_or_path)
+        self.model = AutoModel.from_pretrained(args.model_name_or_path)
+        self.tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path)
 
-        # Define classifier layers
         classifier_dropout = self.args.dropout
         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)
 
-  
-
     def forward(
         self,
-        input_ids=None,
-        attention_mask=None,
-        token_type_ids=None,
-        position_ids=None,
-        signal_bias_mask=None,
-        head_mask=None,
-        inputs_embeds=None,
-        start_positions=None,
-        end_positions=None,
-        output_attentions=None,
-        output_hidden_states=None,
-        return_dict=None,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        token_type_ids: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        signal_bias_mask: Optional[torch.Tensor] = None,
+        head_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        start_positions: Optional[torch.Tensor] = None,     # [batch_size, 3]
+        end_positions: Optional[torch.Tensor] = None,       # [batch_size, 3]
+        output_attentions: Optional[bool] = None,
+        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
         
-        # Ensure that self.model is not None before calling forward
-        if self.model is None:
-            raise ValueError("The model weights have not been loaded. Use from_pretrained() to load them.")
-
-        outputs = self.model(
-            input_ids=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,
-        )
+        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:
+            if self.args.model_name_or_path in ['facebook/bart-large', 'facebook/bart-base', 'facebook/bart-large-cnn']:
+                outputs = self.model(
+                    input_ids,
+                    attention_mask=attention_mask,
+                    head_mask=head_mask,
+                    inputs_embeds=inputs_embeds,
+                    output_attentions=output_attentions,
+                    output_hidden_states=output_hidden_states,
+                    return_dict=return_dict,
+                )
+            elif self.args.model_name_or_path in ['microsoft/deberta-base']:
+                outputs = self.model(
+                    input_ids,
+                    attention_mask=attention_mask,
+                    token_type_ids=token_type_ids,
+                    position_ids=position_ids,
+                    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,
+                )
 
         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)
-
-        # Split logits 
+        logits = self.classifier(sequence_output)  # [batch_size, max_seq_length, 6]
         start_arg0_logits, end_arg0_logits, start_arg1_logits, end_arg1_logits, start_sig_logits, end_sig_logits = logits.split(1, dim=-1)
-        start_arg0_logits = start_arg0_logits.squeeze(-1)
-        end_arg0_logits = end_arg0_logits.squeeze(-1)
-        start_arg1_logits = start_arg1_logits.squeeze(-1)
-        end_arg1_logits = end_arg1_logits.squeeze(-1)
-        start_sig_logits = start_sig_logits.squeeze(-1)
-        end_sig_logits = end_sig_logits.squeeze(-1)
+        start_arg0_logits = start_arg0_logits.squeeze(-1).contiguous()
+        end_arg0_logits = end_arg0_logits.squeeze(-1).contiguous()
+        start_arg1_logits = start_arg1_logits.squeeze(-1).contiguous()
+        end_arg1_logits = end_arg1_logits.squeeze(-1).contiguous()
+        start_sig_logits = start_sig_logits.squeeze(-1).contiguous()
+        end_sig_logits = end_sig_logits.squeeze(-1).contiguous()
+
+        # start_arg0_logits -= (1 - attention_mask) * 1e4
+        # end_arg0_logits -= (1 - attention_mask) * 1e4
+        # start_arg1_logits -= (1 - attention_mask) * 1e4
+        # end_arg1_logits -= (1 - attention_mask) * 1e4
+
+        # start_arg0_logits[:, 0] = -1e4
+        # end_arg0_logits[:, 0] = -1e4
+        # start_arg1_logits[:, 0] = -1e4
+        # end_arg1_logits[:, 0] = -1e4
 
         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, :])
+        # 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 {
             'start_arg0_logits': start_arg0_logits,
@@ -88,27 +208,71 @@ 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
+            '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,
         }
 
-    @classmethod
-    def from_pretrained(cls, model_name, config=None, args=None, **kwargs):
-        """
-        Custom from_pretrained method to load the model from Hugging Face and initialize
-        any additional components such as the classifier.
-        """
-        # Load the configuration
-        config = AutoConfig.from_pretrained(model_name) if config is None else config
-        
-        # Instantiate the model
-        model = cls(args, config)
-        
-        # Load the pre-trained weights into the model
-        model.model = AutoModel.from_pretrained(model_name, config=config, **kwargs, use_safetensors=False)
+    """
+    def save_pretrained(self, save_directory):
         
+        #Save model state dict as safetensor, configuration, and tokenizer files.
         
-        return model
+        # 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,
         start_cause_logits, 
@@ -117,7 +281,7 @@ class ST2ModelV2(nn.Module):
         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
@@ -190,9 +354,27 @@ 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))
@@ -200,12 +382,19 @@ 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()