Add HAT implementation files

modelling_hat.py

@@ -1078,7 +1078,7 @@ class HATForMaskedLM(HATPreTrainedModel):
     def __init__(self, config):
         super().__init__(config)
 
-        self.hat = HATModel(config)
+        self.hi_transformer = HATModel(config)
         self.lm_head = HATLMHead(config)
 
         # The LM head weights require special treatment only when they are tied with the word embeddings
@@ -1123,7 +1123,7 @@ class HATForMaskedLM(HATPreTrainedModel):
         """
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
-        outputs = self.hat(
+        outputs = self.hi_transformer(
             input_ids,
             attention_mask=attention_mask,
             token_type_ids=token_type_ids,
@@ -1161,7 +1161,7 @@ class HATModelForDocumentRepresentation(HATPreTrainedModel):
         self.num_labels = config.num_labels
         self.config = config
 
-        self.hat = HATModel(config)
+        self.hi_transformer = HATModel(config)
         self.pooler = HATPooler(config, pooling=pooling)
 
         # Initialize weights and apply final processing
@@ -1195,7 +1195,7 @@ class HATModelForDocumentRepresentation(HATPreTrainedModel):
         """
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
-        outputs = self.hat(
+        outputs = self.hi_transformer(
             input_ids,
             attention_mask=attention_mask,
             token_type_ids=token_type_ids,
@@ -1237,7 +1237,7 @@ class HATModelForMaskedSentenceRepresentation(HATPreTrainedModel):
         self.num_labels = config.num_labels
         self.config = config
 
-        self.hat = HATModel(config)
+        self.hi_transformer = HATModel(config)
         self.sentencizer = HATSentencizer(config)
 
         # Initialize weights and apply final processing
@@ -1271,7 +1271,7 @@ class HATModelForMaskedSentenceRepresentation(HATPreTrainedModel):
         """
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
-        outputs = self.hat(
+        outputs = self.hi_transformer(
             input_ids,
             attention_mask=attention_mask,
             token_type_ids=token_type_ids,
@@ -1313,7 +1313,7 @@ class HATModelForBoWPreTraining(HATPreTrainedModel):
     def __init__(self, config):
         super().__init__(config)
 
-        self.hat = HATModel(config)
+        self.hi_transformer = HATModel(config)
         if self.config.mlm or self.config.mslm:
             self.lm_head = HATLMHead(config)
         if self.config.srp or self.config.srp:
@@ -1346,7 +1346,7 @@ class HATModelForBoWPreTraining(HATPreTrainedModel):
     ):
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
-        outputs = self.hat(
+        outputs = self.hi_transformer(
             input_ids,
             attention_mask=attention_mask,
             token_type_ids=token_type_ids,
@@ -1447,7 +1447,7 @@ class HATModelForVICRegPreTraining(HATPreTrainedModel):
 
         self.document_regularization = document_regularization
         self.sentence_regularization = sentence_regularization
-        self.hat = HATModel(config)
+        self.hi_transformer = HATModel(config)
         if self.config.mlm:
             self.lm_head = HATLMHead(config)
         if self.config.sent_sim or self.config.doc_sim:
@@ -1474,7 +1474,7 @@ class HATModelForVICRegPreTraining(HATPreTrainedModel):
     ):
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
-        primary_outputs = self.hat(
+        primary_outputs = self.hi_transformer(
             input_ids,
             attention_mask=attention_mask,
             token_type_ids=token_type_ids,
@@ -1484,7 +1484,7 @@ class HATModelForVICRegPreTraining(HATPreTrainedModel):
             return_dict=return_dict,
         )
 
-        secondary_outputs = self.hat(
+        secondary_outputs = self.hi_transformer(
             secondary_input_ids,
             attention_mask=attention_mask,
             token_type_ids=token_type_ids,
@@ -1600,7 +1600,7 @@ class HATModelForSimCLRPreTraining(HATPreTrainedModel):
 
         self.document_regularization = document_regularization
         self.sentence_regularization = sentence_regularization
-        self.hat = HATModel(config)
+        self.hi_transformer = HATModel(config)
         if self.config.mlm:
             self.lm_head = HATLMHead(config)
         if self.config.sent_sim or self.config.doc_sim:
@@ -1626,7 +1626,7 @@ class HATModelForSimCLRPreTraining(HATPreTrainedModel):
     ):
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
-        primary_outputs = self.hat(
+        primary_outputs = self.hi_transformer(
             input_ids,
             attention_mask=attention_mask,
             token_type_ids=token_type_ids,
@@ -1636,7 +1636,7 @@ class HATModelForSimCLRPreTraining(HATPreTrainedModel):
             return_dict=return_dict,
         )
 
-        secondary_outputs = self.hat(
+        secondary_outputs = self.hi_transformer(
             secondary_input_ids,
             attention_mask=attention_mask,
             token_type_ids=token_type_ids,
@@ -1808,7 +1808,7 @@ class HATForSequenceClassification(HATPreTrainedModel):
         self.config = config
         self.pooling = pooling
 
-        self.hat = HATModel(config)
+        self.hi_transformer = HATModel(config)
         classifier_dropout = (
             config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
         )
@@ -1848,7 +1848,7 @@ class HATForSequenceClassification(HATPreTrainedModel):
         """
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
-        outputs = self.hat(
+        outputs = self.hi_transformer(
             input_ids,
             attention_mask=attention_mask,
             token_type_ids=token_type_ids,
@@ -1916,7 +1916,7 @@ class HATModelForSequentialSentenceClassification(HATPreTrainedModel):
         self.num_labels = config.num_labels
         self.config = config
 
-        self.hat = HATModel(config)
+        self.hi_transformer = HATModel(config)
         self.sentencizer = HATSentencizer(config)
         classifier_dropout = (
             config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
@@ -1954,7 +1954,7 @@ class HATModelForSequentialSentenceClassification(HATPreTrainedModel):
         """
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
-        outputs = self.hat(
+        outputs = self.hi_transformer(
             input_ids,
             attention_mask=attention_mask,
             token_type_ids=token_type_ids,
@@ -2020,7 +2020,7 @@ class HATForMultipleChoice(HATPreTrainedModel):
         super().__init__(config)
 
         self.pooling = pooling
-        self.hat = HATModel(config)
+        self.hi_transformer = HATModel(config)
         classifier_dropout = (
             config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
         )
@@ -2071,7 +2071,7 @@ class HATForMultipleChoice(HATPreTrainedModel):
             else None
         )
 
-        outputs = self.hat(
+        outputs = self.hi_transformer(
             flat_input_ids,
             position_ids=flat_position_ids,
             token_type_ids=flat_token_type_ids,
@@ -2125,7 +2125,7 @@ class HATForTokenClassification(HATPreTrainedModel):
         super().__init__(config)
         self.num_labels = config.num_labels
 
-        self.hat = HATModel(config, add_pooling_layer=False)
+        self.hi_transformer = HATModel(config, add_pooling_layer=False)
         classifier_dropout = (
             config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
         )
@@ -2160,7 +2160,7 @@ class HATForTokenClassification(HATPreTrainedModel):
         """
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
-        outputs = self.hat(
+        outputs = self.hi_transformer(
             input_ids,
             attention_mask=attention_mask,
             token_type_ids=token_type_ids,
@@ -2208,7 +2208,7 @@ class HATForQuestionAnswering(HATPreTrainedModel):
         super().__init__(config)
         self.num_labels = config.num_labels
 
-        self.hat = HATModel(config, add_pooling_layer=False)
+        self.hi_transformer = HATModel(config, add_pooling_layer=False)
         self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)
 
         # Initialize weights and apply final processing
@@ -2247,7 +2247,7 @@ class HATForQuestionAnswering(HATPreTrainedModel):
         """
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
-        outputs = self.hat(
+        outputs = self.hi_transformer(
             input_ids,
             attention_mask=attention_mask,
             token_type_ids=token_type_ids,
@@ -2333,3 +2333,4 @@ def off_diagonal(x):
     assert n == m
     return x.flatten()[:-1].view(n - 1, n + 1)[:, 1:].flatten()
 
+