def test_score_examples(self):
example = eval_utils.OpenKpTextExample.from_json(EXAMPLE_JSON)
pr1 = eval_utils.KpPositionPrediction(start_idx=4,
phrase_len=2,
logit=0.1)
pr2 = eval_utils.KpPositionPrediction(start_idx=0,
phrase_len=2,
logit=0.3)
summary = eval_utils.score_examples([example, example], [[pr1], [pr2]])
expected = [0.5, 1 / 6, 1 / 10, 0.25, 0.25, 0.25, 1 / 3, 0.2, 1 / 7]
for i in range(9):
self.assertAlmostEqual(summary[i], expected[i])
def test_get_key_phrase_predictions(self):
example = eval_utils.OpenKpTextExample.from_json(EXAMPLE_JSON)
pr1 = eval_utils.KpPositionPrediction(start_idx=4,
phrase_len=2,
logit=0.1)
pr3 = eval_utils.KpPositionPrediction(start_idx=0,
phrase_len=2,
logit=0.3)
pr2 = eval_utils.KpPositionPrediction(start_idx=2,
phrase_len=2,
logit=0.2)
predictions = example.get_key_phrase_predictions([pr1, pr2, pr3],
max_predictions=2)
self.assertEqual(predictions, ['star trek', 'discovery season'])
def test_get_key_phrase_predictions_skip_invalid_indices(self):
example = eval_utils.OpenKpTextExample.from_json(EXAMPLE_JSON)
pr1 = eval_utils.KpPositionPrediction(start_idx=4,
phrase_len=2,
logit=0.1)
# start_idx=20 is longer than the document.
pr3 = eval_utils.KpPositionPrediction(start_idx=20,
phrase_len=2,
logit=0.3)
pr2 = eval_utils.KpPositionPrediction(start_idx=2,
phrase_len=2,
logit=0.2)
predictions = example.get_key_phrase_predictions([pr1, pr2, pr3],
max_predictions=2)
self.assertEqual(predictions, ['discovery season', '1 director'])
def test_logits_to_predictions(self):
logits = np.array([[0.1, 0.9, 0.5, -0.3], [0.8, 0.3, 0.4, -0.5]])
predictions = eval_utils.logits_to_predictions(logits,
max_predictions=3)
expected1 = eval_utils.KpPositionPrediction(start_idx=1,
phrase_len=1,
logit=0.9)
expected2 = eval_utils.KpPositionPrediction(start_idx=0,
phrase_len=2,
logit=0.8)
expected3 = eval_utils.KpPositionPrediction(start_idx=2,
phrase_len=1,
logit=0.5)
predictions.sort(key=lambda prediction: prediction.logit, reverse=True)
self.assertEqual(predictions[0], expected1)
self.assertEqual(predictions[1], expected2)
self.assertEqual(predictions[2], expected3)