Use the --endpoint flag for reverse sinks as well... add a unit test. Is there a way to check that the unit test is doing what we want?

Author: AngledLuffa

stanza/models/constituency/lstm_model.py

@@ -829,7 +829,10 @@ def map_word(word):
                 else:
                     rattn_inputs = [x + [self.rel_attn_forward(x[0].unsqueeze(0)).squeeze(0)] for x in rattn_inputs]
             if self.rel_attn_reverse is not None:
-                rattn_inputs = [x + [self.rel_attn_reverse(x[0].unsqueeze(0), x[0][-1]).squeeze(0)] for x in rattn_inputs]
+                if self.args['rattn_use_endpoint_sinks']:
+                    rattn_inputs = [x + [self.rel_attn_reverse(x[0].unsqueeze(0), x[0][-1]).squeeze(0)] for x in rattn_inputs]
+                else:
+                    rattn_inputs = [x + [self.rel_attn_reverse(x[0].unsqueeze(0)).squeeze(0)] for x in rattn_inputs]
 
             if self.args['rattn_cat']:
                 all_word_inputs = [torch.cat(x, axis=1) for x in rattn_inputs]

stanza/tests/constituency/test_lstm_model.py

@@ -455,6 +455,16 @@ def test_relative_attention_cat_sinks(pretrain_file):
     model = build_model(pretrain_file, '--no_use_lattn', '--use_rattn', '--rattn_heads', '10', '--rattn_cat', '--rattn_sinks', '2')
     run_forward_checks(model)
 
+def test_relative_attention_endpoint_sinks(pretrain_file):
+    model = build_model(pretrain_file, '--no_use_lattn', '--use_rattn', '--rattn_heads', '10', '--rattn_use_endpoint_sinks', '--rattn_window', '2', '--rattn_sinks', '1')
+    run_forward_checks(model)
+    model = build_model(pretrain_file, '--no_use_lattn', '--use_rattn', '--rattn_heads', '10', '--rattn_use_endpoint_sinks', '--rattn_sinks', '1')
+    run_forward_checks(model)
+    model = build_model(pretrain_file, '--no_use_lattn', '--use_rattn', '--rattn_heads', '10', '--rattn_use_endpoint_sinks', '--rattn_window', '2', '--rattn_sinks', '2')
+    run_forward_checks(model)
+    model = build_model(pretrain_file, '--no_use_lattn', '--use_rattn', '--rattn_heads', '10', '--rattn_use_endpoint_sinks', '--rattn_sinks', '2')
+    run_forward_checks(model)
+
 def test_lstm_tree_forward(pretrain_file):
     """
     Test the LSTM_TREE forward pass