added unzipped gz predictions, the checkpoint with values, and the tree output possibility in the model

checkpoint/checkpoint.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:797ed09acace61d3d397b9d425d28c272f2d6ce8bda8161f6a865fda491526f4
+size 427253662

modeling_ltgbert.py

@@ -374,7 +374,7 @@ class ParserNetwork(nn.Module):
 
         distance, height = self.parse(x, embeddings)
         att_mask, cibling, head, block = self.generate_mask(x, distance, height)
-        return att_mask, cibling, head, block
+        return att_mask, cibling, head, block, distance, height
 
 
 class Encoder(nn.Module):
@@ -790,10 +790,11 @@ LTG_BERT_INPUTS_DOCSTRING = r"""
     LTG_BERT_START_DOCSTRING,
 )
 class LtgBertModel(LtgBertPreTrainedModel):
-    def __init__(self, config, add_mlm_layer=False):
+    def __init__(self, config, add_mlm_layer=False, tree_output=False):
         super().__init__(config)
         self.config = config
-
+        self.tree_output=tree_output
+        
         self.embedding = Embedding(config)
         self.parser_network = ParserNetwork(config, pad=config.pad_token_id)
         self.transformer = Encoder(config, activation_checkpointing=False)
@@ -823,7 +824,7 @@ class LtgBertModel(LtgBertPreTrainedModel):
         device = input_ids.device
 
         static_embeddings, relative_embedding = self.embedding(input_ids.t())
-        att_mask, cibling, head, block = self.parser_network(
+        att_mask, cibling, head, block, distance, height = self.parser_network(
             input_ids.t(), static_embeddings
         )
         contextualized_embeddings, attention_probs = self.transformer(
@@ -837,6 +838,9 @@ class LtgBertModel(LtgBertPreTrainedModel):
             contextualized_embeddings[i] - contextualized_embeddings[i - 1]
             for i in range(1, len(contextualized_embeddings))
         ]
+        if self.tree_output:
+            return last_layer, contextualized_embeddings, attention_probs, {'distance': distance, 'height': height,
+         'cibling': cibling, 'head': head, 'block': block}
         return last_layer, contextualized_embeddings, attention_probs
 
     @add_start_docstrings_to_model_forward(
@@ -863,13 +867,28 @@ class LtgBertModel(LtgBertPreTrainedModel):
         return_dict = (
             return_dict if return_dict is not None else self.config.use_return_dict
         )
-
-        (
-            sequence_output,
-            contextualized_embeddings,
-            attention_probs,
-        ) = self.get_contextualized_embeddings(input_ids, attention_mask)
-
+        tree_values = {} if self.tree_output else None
+        if self.tree_output:
+            (
+                sequence_output,
+                contextualized_embeddings,
+                attention_probs,
+                tree_values
+            ) = self.get_contextualized_embeddings(input_ids, attention_mask)
+        else:
+            (
+                sequence_output,
+                contextualized_embeddings,
+                attention_probs
+            ) = self.get_contextualized_embeddings(input_ids, attention_mask)
+            
+        if self.tree_output:
+            return (
+                sequence_output, 
+                 tree_values,
+                *([contextualized_embeddings] if output_hidden_states else []),
+                *([attention_probs] if output_attentions else []),
+            )
         if not return_dict:
             return (
                 sequence_output,
@@ -891,8 +910,8 @@ class LtgBertModel(LtgBertPreTrainedModel):
 class LtgBertForMaskedLM(LtgBertModel):
     _keys_to_ignore_on_load_unexpected = ["head"]
 
-    def __init__(self, config):
-        super().__init__(config, add_mlm_layer=True)
+    def __init__(self, config, tree_output=False):
+        super().__init__(config, add_mlm_layer=True, tree_output=tree_output)
 
     def get_output_embeddings(self):
         return self.classifier.nonlinearity[-1].weight
@@ -921,12 +940,20 @@ class LtgBertForMaskedLM(LtgBertModel):
         return_dict = (
             return_dict if return_dict is not None else self.config.use_return_dict
         )
-
-        (
-            sequence_output,
-            contextualized_embeddings,
-            attention_probs,
-        ) = self.get_contextualized_embeddings(input_ids, attention_mask)
+        tree_values = {} if self.tree_output else None
+        if self.tree_output:
+            (
+                sequence_output,
+                contextualized_embeddings,
+                attention_probs,
+                tree_values
+            ) = self.get_contextualized_embeddings(input_ids, attention_mask)
+        else:
+            (
+                sequence_output,
+                contextualized_embeddings,
+                attention_probs
+            ) = self.get_contextualized_embeddings(input_ids, attention_mask)
         subword_prediction = self.classifier(sequence_output)
 
         masked_lm_loss = None
@@ -934,7 +961,13 @@ class LtgBertForMaskedLM(LtgBertModel):
             masked_lm_loss = F.cross_entropy(
                 subword_prediction.flatten(0, 1), labels.flatten()
             )
-
+        if self.tree_output:
+            return (
+                sequence_output, 
+                 tree_values,
+                *([contextualized_embeddings] if output_hidden_states else []),
+                *([attention_probs] if output_attentions else []),
+            )
         if not return_dict:
             output = (
                 subword_prediction,
@@ -1027,12 +1060,20 @@ class LtgBertForSequenceClassification(LtgBertModel):
         return_dict = (
             return_dict if return_dict is not None else self.config.use_return_dict
         )
-
-        (
-            sequence_output,
-            contextualized_embeddings,
-            attention_probs,
-        ) = self.get_contextualized_embeddings(input_ids, attention_mask)
+        tree_values = {} if self.tree_output else None
+        if self.tree_output:
+            (
+                sequence_output,
+                contextualized_embeddings,
+                attention_probs,
+                tree_values
+            ) = self.get_contextualized_embeddings(input_ids, attention_mask)
+        else:
+            (
+                sequence_output,
+                contextualized_embeddings,
+                attention_probs
+            ) = self.get_contextualized_embeddings(input_ids, attention_mask)
         logits = self.head(sequence_output[:, 0, :])
 
         loss = None
@@ -1059,7 +1100,14 @@ class LtgBertForSequenceClassification(LtgBertModel):
             elif self.config.problem_type == "multi_label_classification":
                 loss_fct = nn.BCEWithLogitsLoss()
                 loss = loss_fct(logits, labels)
-
+        
+        if self.tree_output:
+            return (
+                sequence_output, 
+                 tree_values,
+                *([contextualized_embeddings] if output_hidden_states else []),
+                *([attention_probs] if output_attentions else []),
+            )
         if not return_dict:
             output = (
                 logits,
@@ -1110,19 +1158,34 @@ class LtgBertForTokenClassification(LtgBertModel):
         return_dict = (
             return_dict if return_dict is not None else self.config.use_return_dict
         )
-
-        (
-            sequence_output,
-            contextualized_embeddings,
-            attention_probs,
-        ) = self.get_contextualized_embeddings(input_ids, attention_mask)
+        tree_values = {} if self.tree_output else None
+        if self.tree_output:
+            (
+                sequence_output,
+                contextualized_embeddings,
+                attention_probs,
+                tree_values
+            ) = self.get_contextualized_embeddings(input_ids, attention_mask)
+        else:
+            (
+                sequence_output,
+                contextualized_embeddings,
+                attention_probs
+            ) = self.get_contextualized_embeddings(input_ids, attention_mask)
         logits = self.head(sequence_output)
 
         loss = None
         if labels is not None:
             loss_fct = nn.CrossEntropyLoss()
             loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
-
+        
+        if self.tree_output:
+            return (
+                sequence_output, 
+                 tree_values,
+                *([contextualized_embeddings] if output_hidden_states else []),
+                *([attention_probs] if output_attentions else []),
+            )
         if not return_dict:
             output = (
                 logits,
@@ -1174,12 +1237,20 @@ class LtgBertForQuestionAnswering(LtgBertModel):
         return_dict = (
             return_dict if return_dict is not None else self.config.use_return_dict
         )
-
-        (
-            sequence_output,
-            contextualized_embeddings,
-            attention_probs,
-        ) = self.get_contextualized_embeddings(input_ids, attention_mask)
+        tree_values = {} if self.tree_output else None
+        if self.tree_output:
+            (
+                sequence_output,
+                contextualized_embeddings,
+                attention_probs,
+                tree_values
+            ) = self.get_contextualized_embeddings(input_ids, attention_mask)
+        else:
+            (
+                sequence_output,
+                contextualized_embeddings,
+                attention_probs
+            ) = self.get_contextualized_embeddings(input_ids, attention_mask)
         logits = self.head(sequence_output)
 
         start_logits, end_logits = logits.split(1, dim=-1)
@@ -1203,7 +1274,14 @@ class LtgBertForQuestionAnswering(LtgBertModel):
             start_loss = loss_fct(start_logits, start_positions)
             end_loss = loss_fct(end_logits, end_positions)
             total_loss = (start_loss + end_loss) / 2
-
+        
+        if self.tree_output:
+            return (
+                sequence_output, 
+                 tree_values,
+                *([contextualized_embeddings] if output_hidden_states else []),
+                *([attention_probs] if output_attentions else []),
+            )
         if not return_dict:
             output = (
                 start_logits,
@@ -1264,12 +1342,20 @@ class LtgBertForMultipleChoice(LtgBertModel):
             if attention_mask is not None
             else None
         )
-
-        (
-            sequence_output,
-            contextualized_embeddings,
-            attention_probs,
-        ) = self.get_contextualized_embeddings(flat_input_ids, flat_attention_mask)
+        tree_values = {} if self.tree_output else None
+        if self.tree_output:
+            (
+                sequence_output,
+                contextualized_embeddings,
+                attention_probs,
+                tree_values
+            ) = self.get_contextualized_embeddings(input_ids, attention_mask)
+        else:
+            (
+                sequence_output,
+                contextualized_embeddings,
+                attention_probs
+            ) = self.get_contextualized_embeddings(input_ids, attention_mask)
         logits = self.head(sequence_output)
         reshaped_logits = logits.view(-1, num_choices)
 
@@ -1277,7 +1363,14 @@ class LtgBertForMultipleChoice(LtgBertModel):
         if labels is not None:
             loss_fct = nn.CrossEntropyLoss()
             loss = loss_fct(reshaped_logits, labels)
-
+        
+        if self.tree_output:
+            return (
+                sequence_output, 
+                 tree_values,
+                *([contextualized_embeddings] if output_hidden_states else []),
+                *([attention_probs] if output_attentions else []),
+            )
         if not return_dict:
             output = (
                 reshaped_logits,