From 3acfb4b3ba36d7d423aabaa95099e426f45ecbcd Mon Sep 17 00:00:00 2001 From: Rogge Niels Date: Tue, 8 Dec 2020 13:00:31 +0100 Subject: [PATCH] Add test for convert_logits_to_predictions method of TapasTokenizer --- tests/test_modeling_tapas.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/tests/test_modeling_tapas.py b/tests/test_modeling_tapas.py index b03fbe98..8f16b076 100644 --- a/tests/test_modeling_tapas.py +++ b/tests/test_modeling_tapas.py @@ -582,6 +582,18 @@ class TapasModelIntegrationTest(unittest.TestCase): self.assertTrue(torch.allclose(logits_aggregation, expected_tensor, atol=TOLERANCE)) + # test the predicted answer coordinates and aggregation indices + EXPECTED_PREDICTED_ANSWER_COORDINATES = [[(0, 0)], [(1, 2)]] + EXPECTED_PREDICTED_AGGREGATION_INDICES = [0, 1] + + predicted_answer_coordinates, predicted_aggregation_indices = tokenizer.convert_logits_to_predictions(inputs, + outputs.logits, + outputs.logits_aggregation) + + self.assertEqual(EXPECTED_PREDICTED_ANSWER_COORDINATES, predicted_answer_coordinates) + self.assertEqual(EXPECTED_PREDICTED_AGGREGATION_INDICES, predicted_aggregation_indices) + + @slow def test_training_question_answering_head_weak_supervision(self): # note that nielsr/tapas-base-finetuned-wtq should correspond to tapas_wtq_wikisql_sqa_inter_masklm_base_reset @@ -685,6 +697,7 @@ class TapasModelIntegrationTest(unittest.TestCase): self.assertTrue(torch.allclose(outputs.logits, expected_tensor, atol=TOLERANCE)) + # Below: tests for Tapas utilities which are defined in modeling_tapas.py. # These are based on segmented_tensor_test.py of the original implementation. # URL: https://github.com/google-research/tapas/blob/master/tapas/models/segmented_tensor_test.py