add targets arg to fill-mask pipeline (#6239)

* add targets arg to fill-mask pipeline

* add tests and more error handling

* quality

* update docstring
This commit is contained in:
Joe Davison
2020-08-12 12:48:29 -04:00
committed by GitHub
parent 0735def8e1
commit bc820476a5
2 changed files with 104 additions and 6 deletions

View File

@@ -41,6 +41,23 @@ expected_fill_mask_result = [
],
]
expected_fill_mask_target_result = [
[
{
"sequence": "<s>My name is Patrick</s>",
"score": 0.004992353264242411,
"token": 3499,
"token_str": "ĠPatrick",
},
{
"sequence": "<s>My name is Clara</s>",
"score": 0.00019297805556561798,
"token": 13606,
"token_str": "ĠClara",
},
]
]
SUMMARIZATION_KWARGS = dict(num_beams=2, min_length=2, max_length=5)
@@ -139,7 +156,7 @@ class MonoColumnInputTestCase(unittest.TestCase):
for key in output_keys:
self.assertIn(key, mono_result[0])
multi_result = [nlp(input) for input in valid_inputs]
multi_result = [nlp(input, **kwargs) for input in valid_inputs]
self.assertIsInstance(multi_result, list)
self.assertIsInstance(multi_result[0], (dict, list))
@@ -219,6 +236,34 @@ class MonoColumnInputTestCase(unittest.TestCase):
nlp, valid_inputs, mandatory_keys, invalid_inputs, expected_check_keys=["sequence"]
)
@require_torch
def test_torch_fill_mask_with_targets(self):
valid_inputs = ["My name is <mask>"]
valid_targets = [[" Teven", " Patrick", " Clara"], [" Sam"]]
invalid_targets = [[], [""], ""]
for model_name in FILL_MASK_FINETUNED_MODELS:
nlp = pipeline(task="fill-mask", model=model_name, tokenizer=model_name, framework="pt")
for targets in valid_targets:
outputs = nlp(valid_inputs, targets=targets)
self.assertIsInstance(outputs, list)
self.assertEqual(len(outputs), len(targets))
for targets in invalid_targets:
self.assertRaises(ValueError, nlp, valid_inputs, targets=targets)
@require_tf
def test_tf_fill_mask_with_targets(self):
valid_inputs = ["My name is <mask>"]
valid_targets = [[" Teven", " Patrick", " Clara"], [" Sam"]]
invalid_targets = [[], [""], ""]
for model_name in FILL_MASK_FINETUNED_MODELS:
nlp = pipeline(task="fill-mask", model=model_name, tokenizer=model_name, framework="tf")
for targets in valid_targets:
outputs = nlp(valid_inputs, targets=targets)
self.assertIsInstance(outputs, list)
self.assertEqual(len(outputs), len(targets))
for targets in invalid_targets:
self.assertRaises(ValueError, nlp, valid_inputs, targets=targets)
@require_torch
@slow
def test_torch_fill_mask_results(self):
@@ -227,6 +272,7 @@ class MonoColumnInputTestCase(unittest.TestCase):
"My name is <mask>",
"The largest city in France is <mask>",
]
valid_targets = [" Patrick", " Clara"]
for model_name in LARGE_FILL_MASK_FINETUNED_MODELS:
nlp = pipeline(task="fill-mask", model=model_name, tokenizer=model_name, framework="pt", topk=2,)
self._test_mono_column_pipeline(
@@ -236,6 +282,14 @@ class MonoColumnInputTestCase(unittest.TestCase):
expected_multi_result=expected_fill_mask_result,
expected_check_keys=["sequence"],
)
self._test_mono_column_pipeline(
nlp,
valid_inputs[:1],
mandatory_keys,
expected_multi_result=expected_fill_mask_target_result,
expected_check_keys=["sequence"],
targets=valid_targets,
)
@require_tf
@slow
@@ -245,6 +299,7 @@ class MonoColumnInputTestCase(unittest.TestCase):
"My name is <mask>",
"The largest city in France is <mask>",
]
valid_targets = [" Patrick", " Clara"]
for model_name in LARGE_FILL_MASK_FINETUNED_MODELS:
nlp = pipeline(task="fill-mask", model=model_name, tokenizer=model_name, framework="tf", topk=2)
self._test_mono_column_pipeline(
@@ -254,6 +309,14 @@ class MonoColumnInputTestCase(unittest.TestCase):
expected_multi_result=expected_fill_mask_result,
expected_check_keys=["sequence"],
)
self._test_mono_column_pipeline(
nlp,
valid_inputs[:1],
mandatory_keys,
expected_multi_result=expected_fill_mask_target_result,
expected_check_keys=["sequence"],
targets=valid_targets,
)
@require_torch
def test_torch_summarization(self):